slab_allocator.hpp

Go to the documentation of this file.

 
#ifndef JOIN_CORE_SLAB_ALLOCATOR_HPP
#define JOIN_CORE_SLAB_ALLOCATOR_HPP
 
// libjoin.
#include <join/backoff.hpp>
#include <join/memory.hpp>
#include <join/utils.hpp>
 
// C++.
#include <algorithm>
#include <utility>
#include <atomic>
#include <array>
#include <tuple>
 
// C.
#include <cassert>
#include <cstdint>
#include <cstddef>
 
namespace join
{
    union alignas (uint64_t) SlabTaggedIndex
    {
        struct
        {
            uint32_t idx;
 
            uint32_t gen;
        };
 
        uint64_t raw;
    };
 
    template <size_t Size>
    union alignas (64) BasicSlabChunk
    {
        static_assert (isPow2 (Size), "size must be a power of 2");
        static_assert (Size % alignof (std::max_align_t) == 0, "size must respects maximum alignment requirement");
        static_assert (Size >= sizeof (uint64_t), "size must respects minimum size for storing index");
 
        std::atomic_uint32_t _next;
 
        uint8_t _data[Size];
    };
 
    template <size_t Size>
    struct BasicSlabSegment
    {
        using Chunk = BasicSlabChunk<Size>;
 
        static constexpr uint64_t MAGIC = 0x9F7E3B2A8D5C4E1B;
 
        static constexpr uint32_t NULL_IDX = UINT32_MAX;
 
        alignas (64) std::atomic_uint64_t _magic;
 
        alignas (64) std::atomic_uint64_t _head;
 
        Chunk _chunks[];
    };
 
    template <size_t Count, size_t Size>
    struct BasicSlabPool
    {
        static_assert (Count > 0, "count must be at least 1");
        static_assert (Count <= UINT32_MAX, "count exceeds tagged index capacity (~256 GB)");
 
        using Chunk = BasicSlabChunk<Size>;
        using Segment = BasicSlabSegment<Size>;
 
        static constexpr size_t _count = Count;
 
        static constexpr size_t _size = Size;
 
        static constexpr size_t _stride = sizeof (Chunk);
 
        static constexpr size_t _total = sizeof (Segment) + _stride * _count;
 
        explicit BasicSlabPool (void* ptr) noexcept
        : _segment (static_cast<Segment*> (ptr))
        {
            uint64_t expected = 0;
 
            if (_segment->_magic.compare_exchange_strong (expected, 0xFFFFFFFFFFFFFFFF, std::memory_order_acq_rel))
            {
                releaseAll ();
 
                _segment->_magic.store (Segment::MAGIC, std::memory_order_release);
            }
            else
            {
                Backoff backoff;
                while (_segment->_magic.load (std::memory_order_acquire) != Segment::MAGIC)
                {
                    backoff ();  // LCOV_EXCL_LINE
                }
            }
        }
 
        BasicSlabPool (const BasicSlabPool& other) = delete;
 
        BasicSlabPool& operator= (const BasicSlabPool& other) = delete;
 
        BasicSlabPool (BasicSlabPool&& other) noexcept
        : _segment (other._segment)
        {
            other._segment = nullptr;
        }
 
        BasicSlabPool& operator= (BasicSlabPool&& other) noexcept
        {
            _segment = other._segment;
            other._segment = nullptr;
            return *this;
        }
 
        ~BasicSlabPool () = default;
 
        void* pop () noexcept
        {
            if (JOIN_UNLIKELY (_segment == nullptr))
            {
                return nullptr;
            }
 
            SlabTaggedIndex cur, next;
            cur.raw = _segment->_head.load (std::memory_order_acquire);
 
            for (;;)
            {
                if (JOIN_UNLIKELY (cur.idx == Segment::NULL_IDX))
                {
                    return nullptr;
                }
 
                next.idx = _segment->_chunks[cur.idx]._next.load (std::memory_order_relaxed);
                next.gen = cur.gen + 1;
 
                if (JOIN_LIKELY (_segment->_head.compare_exchange_weak (cur.raw, next.raw, std::memory_order_acq_rel,
                                                                        std::memory_order_acquire)))
                {
                    return &_segment->_chunks[cur.idx];
                }
            }
        }
 
        void push (void* p) noexcept
        {
            SlabTaggedIndex cur, next;
            next.idx = static_cast<uint32_t> (reinterpret_cast<Chunk*> (p) - _segment->_chunks);
            cur.raw = _segment->_head.load (std::memory_order_relaxed);
 
            for (;;)
            {
                _segment->_chunks[next.idx]._next.store (cur.idx, std::memory_order_relaxed);
                next.gen = cur.gen + 1;
 
                if (JOIN_LIKELY (_segment->_head.compare_exchange_weak (cur.raw, next.raw, std::memory_order_release,
                                                                        std::memory_order_relaxed)))
                {
                    return;
                }
            }
        }
 
        bool reserveAll () noexcept
        {
            SlabTaggedIndex cur, next;
            cur.raw = _segment->_head.load (std::memory_order_acquire);
 
            uint32_t n = 0;
 
            for (uint32_t idx = cur.idx; idx != Segment::NULL_IDX;
                 idx = _segment->_chunks[idx]._next.load (std::memory_order_relaxed))
            {
                ++n;
            }
 
            if (n != _count)
            {
                return false;
            }
 
            next.idx = Segment::NULL_IDX;
            next.gen = cur.gen + 1;
 
            return _segment->_head.compare_exchange_strong (cur.raw, next.raw, std::memory_order_acq_rel,
                                                            std::memory_order_acquire);
        }
 
        void releaseAll () noexcept
        {
            for (uint32_t i = 0; i < _count - 1; ++i)
            {
                _segment->_chunks[i]._next.store (i + 1, std::memory_order_relaxed);
            }
            _segment->_chunks[_count - 1]._next.store (Segment::NULL_IDX, std::memory_order_relaxed);
 
            SlabTaggedIndex cur, next;
            cur.raw = _segment->_head.load (std::memory_order_relaxed);
            next.idx = 0;
            next.gen = cur.gen + 1;
 
            _segment->_head.store (next.raw, std::memory_order_release);
        }
 
        bool owns (void* p) const noexcept
        {
            auto base = reinterpret_cast<std::uintptr_t> (_segment->_chunks);
            auto end = base + (_count * _stride);
            auto ptr = reinterpret_cast<std::uintptr_t> (p);
            return ((ptr >= base) && (ptr < end));
        }
 
        uint32_t getIndex (void* p) const noexcept
        {
            return static_cast<uint32_t> (reinterpret_cast<Chunk*> (p) - _segment->_chunks);
        }
 
        void* getPtr (uint32_t idx) const noexcept
        {
            return &_segment->_chunks[idx];
        }
 
        Segment* _segment = nullptr;
    };
 
    template <typename Backend, size_t Count, size_t... Sizes>
    class BasicSlabAllocator
    {
        static_assert (sizeof...(Sizes) > 0, "slab allocator must have at least one pool size");
        static_assert (std::ranges::is_sorted (std::array{Sizes...}), "pool sizes must be provided in ascending order");
 
    public:
        static constexpr size_t _total = (BasicSlabPool<Count, Sizes>::_total + ...);
 
        template <typename... Args>
        BasicSlabAllocator (Args&&... args)
        : _backend (_total, std::forward<Args> (args)...)
        , _pools (makePools (std::make_index_sequence<sizeof...(Sizes)>{}))
        {
        }
 
        BasicSlabAllocator (const BasicSlabAllocator& other) = delete;
 
        BasicSlabAllocator& operator= (const BasicSlabAllocator& other) = delete;
 
        BasicSlabAllocator (BasicSlabAllocator&& other) noexcept
        : _backend (std::move (other._backend))
        , _pools (std::move (other._pools))
        {
        }
 
        BasicSlabAllocator& operator= (BasicSlabAllocator&& other) noexcept
        {
            _backend = std::move (other._backend);
            _pools = std::move (other._pools);
            return *this;
        }
 
        ~BasicSlabAllocator () = default;
 
        void* allocate (size_t size) noexcept
        {
            return allocateImplem<0> (size);
        }
 
        void* tryAllocate (size_t size) noexcept
        {
            return tryAllocateImplem<0> (size);
        }
 
        void deallocate (void* p) noexcept
        {
            if (p == nullptr)
            {
                return;
            }
            deallocateImplem<0> (p);
        }
 
        void deallocate (void* p, [[maybe_unused]] size_t size) noexcept
        {
            deallocate (p);
        }
 
        bool hasBackend () const noexcept
        {
            return _backend.mapped ();
        }
 
        template <size_t I = 0>
        uint32_t getIndex (void* p) const noexcept
        {
            static_assert (I < sizeof...(Sizes), "pool index out of range");
            return std::get<I> (_pools).getIndex (p);
        }
 
        template <size_t I = 0>
        void* getPtr (uint32_t idx) const noexcept
        {
            static_assert (I < sizeof...(Sizes), "pool index out of range");
            return std::get<I> (_pools).getPtr (idx);
        }
 
        template <size_t I = 0>
        bool reserveAll () noexcept
        {
            static_assert (I < sizeof...(Sizes), "pool index out of range");
            return std::get<I> (_pools).reserveAll ();
        }
 
        template <size_t I = 0>
        void releaseAll () noexcept
        {
            static_assert (I < sizeof...(Sizes), "pool index out of range");
            std::get<I> (_pools).releaseAll ();
        }
 
#ifdef JOIN_HAS_NUMA
        int mbind (int numa) const noexcept
        {
            return _backend.mbind (numa);
        }
#endif
 
        int mlock () const noexcept
        {
            return _backend.mlock ();
        }
 
    private:
        static std::array<size_t, sizeof...(Sizes)> makeOffsets ()
        {
            size_t totals[] = {BasicSlabPool<Count, Sizes>::_total...};
            std::array<size_t, sizeof...(Sizes)> offsets{};
            for (size_t i = 1; i < sizeof...(Sizes); ++i)
            {
                offsets[i] = offsets[i - 1] + totals[i - 1];
            }
            return offsets;
        }
 
        template <size_t... Is>
        std::tuple<BasicSlabPool<Count, Sizes>...> makePools (std::index_sequence<Is...>) noexcept
        {
            auto offsets = makeOffsets ();
            return std::tuple<BasicSlabPool<Count, Sizes>...>{
                BasicSlabPool<Count, Sizes> (static_cast<char*> (_backend.get ()) + offsets[Is])...};
        }
 
        template <size_t I>
        void* allocateImplem (size_t size) noexcept
        {
            if constexpr (I < sizeof...(Sizes))
            {
                auto& pool = std::get<I> (_pools);
                if (size <= pool._size)
                {
                    void* p = pool.pop ();
                    if (p != nullptr)
                    {
                        return p;
                    }
                }
                return allocateImplem<I + 1> (size);
            }
            else
            {
                return nullptr;
            }
        }
 
        template <size_t I>
        void* tryAllocateImplem (size_t size) noexcept
        {
            if constexpr (I < sizeof...(Sizes))
            {
                auto& pool = std::get<I> (_pools);
                if (size <= pool._size)
                {
                    return pool.pop ();
                }
                return tryAllocateImplem<I + 1> (size);
            }
            else
            {
                return nullptr;
            }
        }
 
        template <size_t I>
        void deallocateImplem (void* p) noexcept
        {
            if constexpr (I < sizeof...(Sizes))
            {
                auto& pool = std::get<I> (_pools);
                if (pool.owns (p))
                {
                    pool.push (p);
                    return;
                }
                deallocateImplem<I + 1> (p);
            }
        }
 
        Backend _backend;
 
        std::tuple<BasicSlabPool<Count, Sizes>...> _pools;
    };
}
 
#endif