ring.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "condy/condy_uring.hpp"
#include <cassert>
#include <cerrno>
#include <cstddef>
#include <cstring>
 
namespace condy {

class FdTable {
public:
    FdTable(io_uring &ring) : ring_(ring) {}
 
    FdTable(const FdTable &) = delete;
    FdTable &operator=(const FdTable &) = delete;
    FdTable(FdTable &&) = delete;
    FdTable &operator=(FdTable &&) = delete;
 
public:
    int init(size_t capacity) noexcept {
        return io_uring_register_files_sparse(&ring_, capacity);
    }

    int init(const int *fds, unsigned nr_fds) {
        return io_uring_register_files(&ring_, fds, nr_fds);
    }

    int destroy() noexcept { return io_uring_unregister_files(&ring_); }

    int update(unsigned index_base, const int *fds, unsigned nr_fds) noexcept {
        return io_uring_register_files_update(&ring_, index_base, fds, nr_fds);
    }

    int set_file_alloc_range(unsigned offset, unsigned size) noexcept {
        return io_uring_register_file_alloc_range(&ring_, offset, size);
    }
 
private:
    io_uring &ring_;
 
    friend class Runtime;
};

class BufferTable {
public:
    BufferTable(io_uring &ring) : ring_(ring) {}
 
    BufferTable(const BufferTable &) = delete;
    BufferTable &operator=(const BufferTable &) = delete;
    BufferTable(BufferTable &&) = delete;
    BufferTable &operator=(BufferTable &&) = delete;
 
public:
    int init(size_t capacity) noexcept {
        int r = io_uring_register_buffers_sparse(&ring_, capacity);
        if (r < 0) {
            return r;
        }
        initialized_ = true;
        return r;
    }

    int init(const iovec *vecs, unsigned nr_vecs) {
        int r = io_uring_register_buffers(&ring_, vecs, nr_vecs);
        if (r < 0) {
            return r;
        }
        initialized_ = true;
        return r;
    }

    int destroy() noexcept {
        initialized_ = false;
        return io_uring_unregister_buffers(&ring_);
    }

    int update(unsigned index_base, const iovec *vecs,
               unsigned nr_vecs) noexcept {
        return io_uring_register_buffers_update_tag(&ring_, index_base, vecs,
                                                    nullptr, nr_vecs);
    }
 
#if !IO_URING_CHECK_VERSION(2, 10) // >= 2.10

    int clone_buffers(BufferTable &src, unsigned int dst_off = 0,
                      unsigned int src_off = 0, unsigned int nr = 0) noexcept {
        auto *src_ring = &src.ring_;
        auto *dst_ring = &ring_;
        unsigned int flags = 0;
        if (initialized_) {
            flags |= IORING_REGISTER_DST_REPLACE;
        }
        int r = __io_uring_clone_buffers_offset(dst_ring, src_ring, dst_off,
                                                src_off, nr, flags);
        if (r < 0) {
            return r;
        }
        initialized_ = true;
        return r;
    }
#endif
 
private:
    io_uring &ring_;
    bool initialized_ = false;
};

class RingSettings {
public:
    RingSettings(io_uring &ring) : ring_(ring) {}
 
    ~RingSettings() {
        if (probe_) {
            io_uring_free_probe(probe_);
            probe_ = nullptr;
        }
    }
 
    RingSettings(const RingSettings &) = delete;
    RingSettings &operator=(const RingSettings &) = delete;
    RingSettings(RingSettings &&) = delete;
    RingSettings &operator=(RingSettings &&) = delete;
 
public:
    int apply_iowq_aff(size_t cpusz, const cpu_set_t *mask) noexcept {
        return io_uring_register_iowq_aff(&ring_, cpusz, mask);
    }

    int remove_iowq_aff() noexcept {
        return io_uring_unregister_iowq_aff(&ring_);
    }

    int set_iowq_max_workers(unsigned int *values) noexcept {
        return io_uring_register_iowq_max_workers(&ring_, values);
    }

    io_uring_probe *get_probe() noexcept {
        if (probe_) {
            return probe_;
        }
        probe_ = io_uring_get_probe_ring(&ring_);
        return probe_;
    }

    uint32_t get_features() const noexcept { return features_; }
 
#if !IO_URING_CHECK_VERSION(2, 6) // >= 2.6
    int apply_napi(io_uring_napi *napi) noexcept {
        return io_uring_register_napi(&ring_, napi);
    }

    int remove_napi(io_uring_napi *napi = nullptr) noexcept {
        return io_uring_unregister_napi(&ring_, napi);
    }
#endif
 
#if !IO_URING_CHECK_VERSION(2, 8) // >= 2.8
    int set_clock(io_uring_clock_register *clock_reg) noexcept {
        return io_uring_register_clock(&ring_, clock_reg);
    }
#endif
 
#if !IO_URING_CHECK_VERSION(2, 9) // >= 2.9
    int set_rings_size(io_uring_params *params) noexcept {
        return io_uring_resize_rings(&ring_, params);
    }
#endif
 
#if !IO_URING_CHECK_VERSION(2, 10) // >= 2.10
    int set_iowait(bool enable_iowait) noexcept {
        return io_uring_set_iowait(&ring_, enable_iowait);
    }
#endif
 
private:
    io_uring &ring_;
    io_uring_probe *probe_ = nullptr;
    uint32_t features_ = 0;
 
    friend class Ring;
};
 
class Ring {
public:
    Ring() = default;
    ~Ring() { destroy(); }
 
    Ring(const Ring &) = delete;
    Ring &operator=(const Ring &) = delete;
    Ring(Ring &&) = delete;
    Ring &operator=(Ring &&) = delete;
 
public:
    int init(unsigned int entries, io_uring_params *params,
             [[maybe_unused]] void *buf = nullptr,
             [[maybe_unused]] size_t buf_size = 0) noexcept {
        int r;
        if (initialized_) {
            return -EBUSY;
        }
#if !IO_URING_CHECK_VERSION(2, 5) // >= 2.5
        if (params->flags & IORING_SETUP_NO_MMAP) {
            r = io_uring_queue_init_mem(entries, &ring_, params, buf, buf_size);
        } else
#endif
            r = io_uring_queue_init_params(entries, &ring_, params);
        if (r < 0) {
            return r;
        }
        settings_.features_ = params->features;
        sqpoll_mode_ = (params->flags & IORING_SETUP_SQPOLL) != 0;
        initialized_ = true;
        return r;
    }
 
    void destroy() noexcept {
        if (initialized_) {
            io_uring_queue_exit(&ring_);
            initialized_ = false;
        }
    }
 
    void submit() noexcept { io_uring_submit(&ring_); }
 
    template <typename Func>
    ssize_t reap_completions_wait(Func &&process_func) noexcept {
        unsigned head;
        io_uring_cqe *cqe;
        ssize_t reaped = 0;
        do {
            int r = io_uring_submit_and_wait(&ring_, 1);
            if (r >= 0) [[likely]] {
                break;
            } else if (r == -EINTR) {
                continue;
            } else {
                return r;
            }
        } while (true);
 
        io_uring_for_each_cqe(&ring_, head, cqe) {
            process_func(cqe);
#if !IO_URING_CHECK_VERSION(2, 13) // >= 2.13
            reaped += io_uring_cqe_nr(cqe);
#else
            reaped++;
#endif
        }
        io_uring_cq_advance(&ring_, reaped);
        return reaped;
    }
 
    template <typename Func>
    ssize_t reap_completions(Func &&process_func) noexcept {
        io_uring_cqe *cqe;
        int r = io_uring_peek_cqe(&ring_, &cqe);
        if (r == -EAGAIN) {
            return 0;
        } else if (r < 0) {
            return r;
        }
 
        unsigned head;
        ssize_t reaped = 0;
        io_uring_for_each_cqe(&ring_, head, cqe) {
            process_func(cqe);
#if !IO_URING_CHECK_VERSION(2, 13) // >= 2.13
            reaped += io_uring_cqe_nr(cqe);
#else
            reaped++;
#endif
        }
        io_uring_cq_advance(&ring_, reaped);
        return reaped;
    }
 
    void reserve_space(size_t n) noexcept {
        size_t space_left;
        do {
            space_left = io_uring_sq_space_left(&ring_);
            if (space_left >= n) {
                return;
            }
            submit();
        } while (true);
    }
 
    io_uring *ring() noexcept { return &ring_; }
 
    FdTable &fd_table() noexcept { return fd_table_; }
 
    BufferTable &buffer_table() noexcept { return buffer_table_; }
 
    RingSettings &settings() noexcept { return settings_; }
 
    io_uring_sqe *get_sqe() noexcept { return get_sqe_<io_uring_get_sqe>(); }
 
#if !IO_URING_CHECK_VERSION(2, 13) // >= 2.13
    io_uring_sqe *get_sqe128() noexcept {
        if (ring_.flags & (IORING_SETUP_SQE128 | IORING_SETUP_SQE_MIXED))
            [[likely]] {
            return get_sqe_<io_uring_get_sqe128>();
        }
        return nullptr;
    }
#endif
 
private:
    template <io_uring_sqe *(*get_sqe)(struct io_uring *)>
    io_uring_sqe *get_sqe_() noexcept {
        [[maybe_unused]] int r;
        io_uring_sqe *sqe;
        do {
            sqe = get_sqe(&ring_);
            if (sqe) {
                break;
            }
            r = io_uring_submit(&ring_);
            assert(r >= 0);
            if (sqpoll_mode_) {
                r = io_uring_sqring_wait(&ring_);
                assert(r >= 0);
            }
        } while (true);
        return sqe;
    }
 
private:
    bool initialized_ = false;
    io_uring ring_;
    bool sqpoll_mode_ = false;
 
    FdTable fd_table_{ring_};
    BufferTable buffer_table_{ring_};
    RingSettings settings_{ring_};
};
 
} // namespace condy