provided_buffers.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "condy/concepts.hpp"
#include "condy/condy_uring.hpp"
#include "condy/detail/buffers.hpp"
#include "condy/detail/context.hpp"
#include "condy/detail/ring.hpp"
#include "condy/detail/utils.hpp"
#include "condy/runtime.hpp"
#include <algorithm>
#include <bit>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <sys/mman.h>
#include <sys/types.h>
#include <vector>
 
namespace condy {

struct BufferInfo {
    uint16_t bid;
    uint16_t num_buffers;
};
 
namespace detail {
 
class BundledProvidedBufferQueue {
protected:
    BundledProvidedBufferQueue(Runtime *runtime, uint32_t capacity,
                               unsigned int flags)
        : runtime_(runtime), capacity_(std::bit_ceil(capacity)),
          mask_(io_uring_buf_ring_mask(capacity_)), buf_lens_(capacity_, 0),
          br_flags_(flags) {
        bool ok = false;
        auto d = detail::defer([&]() {
            if (!ok) {
                cleanup_();
            }
        });
 
        auto &bgid_pool = runtime_->bgid_pool_internal();
 
        bgid_ = bgid_pool.allocate();
 
        int err = 0;
        br_ = io_uring_setup_buf_ring(runtime_->ring_internal().ring(),
                                      capacity_, bgid_, br_flags_, &err);
        if (br_ == nullptr) [[unlikely]] {
            throw detail::make_system_error("io_uring_setup_buf_ring", -err);
        }
 
        ok = true;
    }
 
    ~BundledProvidedBufferQueue() { cleanup_(); }
 
public:
    CONDY_DELETE_COPY_MOVE(BundledProvidedBufferQueue);
 
public:
    size_t size() const noexcept { return size_; }

    size_t capacity() const noexcept { return capacity_; }

    template <BufferLike Buffer> uint16_t push(const Buffer &buffer) {
        if (size_ >= capacity_) [[unlikely]] {
            throw std::logic_error("Capacity exceeded");
        }
 
        uint16_t bid = br_->tail & mask_;
        io_uring_buf_ring_add(br_, buffer.data(), buffer.size(), bid, mask_, 0);
        buf_lens_[bid] = buffer.size();
        io_uring_buf_ring_advance(br_, 1);
        size_++;
 
        return bid;
    }
 
public:
    uint16_t bgid() const noexcept { return bgid_; }
 
    BufferInfo handle_finish(io_uring_cqe *cqe) noexcept {
        assert(cqe != nullptr);
        int32_t res = cqe->res;
        uint32_t flags = cqe->flags;
 
        if (!(flags & IORING_CQE_F_BUFFER)) {
            return BufferInfo{0, 0};
        }
 
        assert(res > 0);
 
        BufferInfo result = {
            .bid = static_cast<uint16_t>(flags >> IORING_CQE_BUFFER_SHIFT),
            .num_buffers = 0,
        };
 
        bool is_incr = false;
#if CONDY_URING_VERSION_GE(2, 8) // >= 2.8
        is_incr = br_flags_ & IOU_PBUF_RING_INC;
#endif
 
        uint16_t idx = result.bid;
        int64_t bytes = res;
        while (bytes > 0) {
            uint16_t curr_bid = idx & mask_;
            uint32_t buf_len;
            if (is_incr) {
                buf_len = std::min<uint32_t>(bytes, buf_lens_[curr_bid]);
                buf_lens_[curr_bid] -= buf_len;
            } else {
                buf_len = std::exchange(buf_lens_[curr_bid], 0);
            }
            bytes -= buf_len;
            if (buf_lens_[curr_bid] == 0) {
                result.num_buffers++;
            }
            idx++;
        }
        assert(size_ >= result.num_buffers);
        size_ -= result.num_buffers;
 
        return result;
    }
 
private:
    void cleanup_() noexcept {
        if (br_ != nullptr) {
            [[maybe_unused]] int r = io_uring_free_buf_ring(
                runtime_->ring_internal().ring(), br_, capacity_, bgid_);
            assert(r == 0);
            br_ = nullptr;
        }
        if (bgid_ != std::numeric_limits<uint16_t>::max()) {
            runtime_->bgid_pool_internal().recycle(bgid_);
        }
    }
 
private:
    Runtime *runtime_;
    io_uring_buf_ring *br_ = nullptr;
    uint32_t size_ = 0;
    uint32_t capacity_;
    int mask_;
    uint16_t bgid_ = std::numeric_limits<uint16_t>::max();
    std::vector<uint32_t> buf_lens_;
    unsigned int br_flags_;
};
 
} // namespace detail

class ProvidedBufferQueue : public detail::BundledProvidedBufferQueue {
public:
    ProvidedBufferQueue(uint32_t capacity, unsigned int flags = 0)
        : ProvidedBufferQueue(*detail::Context::current().runtime(), capacity,
                              flags) {}

    ProvidedBufferQueue(Runtime &runtime, uint32_t capacity,
                        unsigned int flags = 0)
        : BundledProvidedBufferQueue(&runtime, capacity, flags) {}
 
    BufferInfo handle_finish(io_uring_cqe *cqe) noexcept {
        assert(cqe != nullptr);
        auto result = BundledProvidedBufferQueue::handle_finish(cqe);
        assert(result.num_buffers <= 1);
        return result;
    }
};
 
namespace detail {
class BundledProvidedBufferPool;
}

class ProvidedBuffer
    : public detail::ManagedBuffer<detail::BundledProvidedBufferPool> {
public:
    using Base = detail::ManagedBuffer<detail::BundledProvidedBufferPool>;
    using Base::Base;
};
 
namespace detail {
 
class BundledProvidedBufferPool {
protected:
    BundledProvidedBufferPool(Runtime *runtime, void *buffer_data,
                              uint32_t num_buffers, size_t buffer_size,
                              unsigned int flags, bool owns_buffers)
        : runtime_(runtime), buffers_base_(static_cast<char *>(buffer_data)),
          num_buffers_(std::bit_ceil(num_buffers)),
          mask_(io_uring_buf_ring_mask(num_buffers_)),
          buffer_size_(buffer_size), curr_buf_len_(buffer_size),
          br_flags_(flags), owns_buffers_(owns_buffers) {
        bool ok = false;
        auto d = detail::defer([&]() {
            if (!ok) {
                cleanup_();
            }
        });
 
        auto &bgid_pool = runtime_->bgid_pool_internal();
 
        bgid_ = bgid_pool.allocate();
 
        int err = 0;
        br_ = io_uring_setup_buf_ring(runtime_->ring_internal().ring(),
                                      num_buffers_, bgid_, br_flags_, &err);
        if (br_ == nullptr) [[unlikely]] {
            throw detail::make_system_error("io_uring_setup_buf_ring", -err);
        }
 
        for (size_t bid = 0; bid < num_buffers_; bid++) {
            char *ptr = buffers_base_ + bid * buffer_size_;
            io_uring_buf_ring_add(br_, ptr, buffer_size_, bid, mask_,
                                  static_cast<int>(bid));
        }
        io_uring_buf_ring_advance(br_, static_cast<int>(num_buffers_));
 
        ok = true;
    }
 
    ~BundledProvidedBufferPool() { cleanup_(); }
 
public:
    CONDY_DELETE_COPY_MOVE(BundledProvidedBufferPool);
 
public:
    size_t capacity() const noexcept { return num_buffers_; }

    size_t buffer_size() const noexcept { return buffer_size_; }
 
public:
    uint16_t bgid() const noexcept { return bgid_; }
 
    std::vector<ProvidedBuffer> handle_finish(io_uring_cqe *cqe) noexcept {
        assert(cqe != nullptr);
        int32_t res = cqe->res;
        uint32_t flags = cqe->flags;
        std::vector<ProvidedBuffer> buffers;
 
        if (!(flags & IORING_CQE_F_BUFFER)) {
            return buffers;
        }
 
        assert(res > 0);
 
        bool is_incr = false;
#if CONDY_URING_VERSION_GE(2, 8) // >= 2.8
        is_incr = br_flags_ & IOU_PBUF_RING_INC;
#endif
 
        int64_t bytes = res;
        while (bytes > 0) {
            auto *buf_ptr = curr_io_uring_buf_();
            uint16_t bid = buf_ptr->bid;
 
            char *data = get_buffer_(bid) + (buffer_size_ - curr_buf_len_);
            uint32_t buf_len;
            if (is_incr) {
                buf_len = std::min<uint32_t>(bytes, curr_buf_len_);
                curr_buf_len_ -= buf_len;
            } else {
                buf_len = std::exchange(curr_buf_len_, 0);
            }
            bytes -= buf_len;
            if (curr_buf_len_ == 0) {
                buffers.emplace_back(data, buf_len, this);
                advance_io_uring_buf_();
                curr_buf_len_ = buffer_size_;
            } else {
                buffers.emplace_back(data, buf_len, nullptr);
            }
        }
 
        return buffers;
    }
 
    void add_buffer_back(void *ptr, [[maybe_unused]] size_t size) noexcept {
        assert(size <= buffer_size_);
        assert(ptr >= buffers_base_);
        size_t offset = static_cast<char *>(ptr) - buffers_base_;
        size_t bid = offset / buffer_size_;
        assert(bid < num_buffers_);
        char *buffer_ptr = buffers_base_ + bid * buffer_size_;
        io_uring_buf_ring_add(br_, buffer_ptr, buffer_size_, bid, mask_, 0);
        io_uring_buf_ring_advance(br_, 1);
    }
 
private:
    char *get_buffer_(uint16_t bid) const noexcept {
        return buffers_base_ + static_cast<size_t>(bid) * buffer_size_;
    }
 
    io_uring_buf *curr_io_uring_buf_() noexcept {
        return &br_->bufs[br_head_ & mask_];
    }
 
    void advance_io_uring_buf_() noexcept { br_head_++; }
 
    void cleanup_() noexcept {
        [[maybe_unused]] int r;
        if (br_ != nullptr) {
            r = io_uring_free_buf_ring(runtime_->ring_internal().ring(), br_,
                                       num_buffers_, bgid_);
            assert(r == 0);
            br_ = nullptr;
        }
        if (bgid_ != std::numeric_limits<uint16_t>::max()) {
            runtime_->bgid_pool_internal().recycle(bgid_);
        }
        if (owns_buffers_) {
            r = munmap(buffers_base_,
                       static_cast<size_t>(num_buffers_) * buffer_size_);
            assert(r == 0);
            owns_buffers_ = false;
            buffers_base_ = nullptr;
        }
    }
 
private:
    Runtime *runtime_;
    io_uring_buf_ring *br_ = nullptr;
    char *buffers_base_ = nullptr;
    uint32_t num_buffers_;
    int mask_;
    uint32_t buffer_size_;
    uint32_t curr_buf_len_;
    uint16_t bgid_ = std::numeric_limits<uint16_t>::max();
    uint16_t br_head_ = 0;
    unsigned int br_flags_;
    bool owns_buffers_;
};
 
} // namespace detail

class ProvidedBufferPool : public detail::BundledProvidedBufferPool {
public:
    ProvidedBufferPool(uint32_t num_buffers, size_t buffer_size,
                       unsigned int flags = 0)
        : ProvidedBufferPool(*detail::Context::current().runtime(), num_buffers,
                             buffer_size, flags) {}

    ProvidedBufferPool(Runtime &runtime, uint32_t num_buffers,
                       size_t buffer_size, unsigned int flags = 0)
        : BundledProvidedBufferPool(
              &runtime,
              alloc_buffer_data_(std::bit_ceil(num_buffers) * buffer_size),
              num_buffers, buffer_size, flags, true) {}

    ProvidedBufferPool(void *buffer_data, uint32_t num_buffers,
                       size_t buffer_size, unsigned int flags = 0)
        : ProvidedBufferPool(*detail::Context::current().runtime(), buffer_data,
                             num_buffers, buffer_size, flags) {}

    ProvidedBufferPool(Runtime &runtime, void *buffer_data,
                       uint32_t num_buffers, size_t buffer_size,
                       unsigned int flags = 0)
        : BundledProvidedBufferPool(&runtime, buffer_data, num_buffers,
                                    buffer_size, flags, false) {}
 
public:
    ProvidedBuffer handle_finish(io_uring_cqe *cqe) noexcept {
        assert(cqe != nullptr);
        auto buffers = BundledProvidedBufferPool::handle_finish(cqe);
        if (buffers.empty()) {
            return ProvidedBuffer();
        }
        assert(buffers.size() == 1);
        return std::move(buffers[0]);
    }
 
private:
    static void *alloc_buffer_data_(size_t buf_size) {
        void *buf_data = mmap(nullptr, buf_size, PROT_READ | PROT_WRITE,
                              MAP_ANONYMOUS | MAP_PRIVATE, 0, 0);
        if (buf_data == MAP_FAILED) [[unlikely]] {
            throw detail::make_system_error("mmap");
        }
        return buf_data;
    }
};

inline auto &bundled(ProvidedBufferPool &buffer) {
    return static_cast<detail::BundledProvidedBufferPool &>(buffer);
}

inline auto &bundled(ProvidedBufferQueue &buffer) {
    return static_cast<detail::BundledProvidedBufferQueue &>(buffer);
}
 
} // namespace condy