allocator.hpp

Go to the documentation of this file.

 
#ifndef JOIN_CORE_ALLOCATOR_HPP
#define JOIN_CORE_ALLOCATOR_HPP
 
// libjoin.
#include <join/backoff.hpp>
#include <join/memory.hpp>
#include <join/utils.hpp>
 
// C++.
#include <utility>
#include <atomic>
#include <array>
#include <tuple>
 
// C.
#include <cassert>
#include <cstdint>
#include <cstddef>
 
namespace join
{
    union alignas (uint64_t) TaggedIndex
    {
        struct
        {
            uint32_t idx;
 
            uint32_t gen;
        };
 
        uint64_t raw;
    };
 
    template <size_t Size>
    union alignas (std::max_align_t) BasicChunk
    {
        static_assert ((Size & (Size - 1)) == 0, "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");
 
        uint32_t _next;
 
        uint8_t _data[Size];
    };
 
    template <size_t Size>
    struct BasicSegment
    {
        using Chunk = BasicChunk<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 BasicPool
    {
        static_assert (Count > 0, "count must be at least 1");
        static_assert (Count <= UINT32_MAX, "count exceeds tagged index capacity (~256 GB)");
 
        using Chunk   = BasicChunk<Size>;
        using Segment = BasicSegment<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 BasicPool (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))
            {
                for (uint32_t i = 0; i < _count - 1; ++i)
                {
                    _segment->_chunks[i]._next = i + 1;
                }
                _segment->_chunks[_count - 1]._next = Segment::NULL_IDX;
 
                TaggedIndex head;
                head.idx = 0;
                head.gen = 0;
                _segment->_head.store (head.raw, std::memory_order_relaxed);
 
                _segment->_magic.store (Segment::MAGIC, std::memory_order_release);
            }
            else
            {
                Backoff backoff;
                while (_segment->_magic.load (std::memory_order_acquire) != Segment::MAGIC)
                {
                    backoff ();
                }
            }
        }
 
        BasicPool (const BasicPool& other) = delete;
 
        BasicPool& operator= (const BasicPool& other) = delete;
 
        BasicPool (BasicPool&& other) noexcept
        : _segment (other._segment)
        {
            other._segment = nullptr;
        }
 
        BasicPool& operator= (BasicPool&& other) noexcept
        {
            _segment       = other._segment;
            other._segment = nullptr;
            return *this;
        }
 
        ~BasicPool () = default;
 
        void* pop () noexcept
        {
            TaggedIndex 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;
                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
        {
            TaggedIndex 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 = cur.idx;
                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 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));
        }
 
        Segment* _segment = nullptr;
    };
 
    template <size_t Count, size_t... Sizes>
    struct TotalSize;
 
    template <size_t Count, size_t First, size_t... Rest>
    struct TotalSize<Count, First, Rest...>
    {
        static constexpr size_t value = BasicPool<Count, First>::_total + TotalSize<Count, Rest...>::value;
    };
 
    template <size_t Count, size_t Last>
    struct TotalSize<Count, Last>
    {
        static constexpr size_t value = BasicPool<Count, Last>::_total;
    };
 
    template <size_t... Sizes>
    struct Sorted;
 
    template <size_t First, size_t Second, size_t... Rest>
    struct Sorted<First, Second, Rest...>
    : std::integral_constant<bool, (First <= Second) && Sorted<Second, Rest...>::value>
    {
    };
 
    template <size_t Last>
    struct Sorted<Last> : std::true_type
    {
    };
 
    template <>
    struct Sorted<> : std::true_type
    {
    };
 
    template <typename Backend, size_t Count, size_t... Sizes>
    class BasicArena
    {
        static_assert (sizeof...(Sizes) > 0, "arena must have at least one pool size");
        static_assert (Sorted<Sizes...>::value, "pool sizes must be provided in ascending order");
 
    public:
        static constexpr size_t _total = TotalSize<Count, Sizes...>::value;
 
        template <typename... Args>
        BasicArena (Args&&... args)
        : _backend (_total, std::forward<Args> (args)...)
        , _pools (makePools (std::make_index_sequence<sizeof...(Sizes)>{}))
        {
        }
 
        BasicArena (const BasicArena& other) = delete;
 
        BasicArena& operator= (const BasicArena& other) = delete;
 
        BasicArena (BasicArena&& other) noexcept
        : _backend (std::move (other._backend))
        , _pools (std::move (other._pools))
        {
        }
 
        BasicArena& operator= (BasicArena&& other) noexcept
        {
            _backend = std::move (other._backend);
            _pools   = std::move (other._pools);
            return *this;
        }
 
        ~BasicArena () = 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);
        }
 
        int mbind (int numa) const noexcept
        {
            return _backend.mbind (numa);
        }
 
        int mlock () const noexcept
        {
            return _backend.mlock ();
        }
 
    private:
        static std::array<size_t, sizeof...(Sizes)> makeOffsets ()
        {
            size_t totals[] = {BasicPool<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<BasicPool<Count, Sizes>...> makePools (std::index_sequence<Is...>) noexcept
        {
            auto offsets = makeOffsets ();
            return std::tuple<BasicPool<Count, Sizes>...>{
                BasicPool<Count, Sizes> (static_cast<char*> (_backend.get ()) + offsets[Is])...};
        }
 
        template <size_t I>
        typename std::enable_if<(I < sizeof...(Sizes)), void*>::type allocateImplem (size_t size) noexcept
        {
            auto& pool = std::get<I> (_pools);
            if (size <= pool._size)
            {
                void* p = pool.pop ();
                if (p != nullptr)
                {
                    return p;
                }
            }
            return allocateImplem<I + 1> (size);
        }
 
        template <size_t I>
        typename std::enable_if<(I >= sizeof...(Sizes)), void*>::type allocateImplem (size_t) noexcept
        {
            return nullptr;
        }
 
        template <size_t I>
        typename std::enable_if<(I < sizeof...(Sizes)), void*>::type tryAllocateImplem (size_t size) noexcept
        {
            auto& pool = std::get<I> (_pools);
            if (size <= pool._size)
            {
                return pool.pop ();
            }
            return tryAllocateImplem<I + 1> (size);
        }
 
        template <size_t I>
        typename std::enable_if<(I >= sizeof...(Sizes)), void*>::type tryAllocateImplem (size_t) noexcept
        {
            return nullptr;
        }
 
        template <size_t I>
        typename std::enable_if<(I < sizeof...(Sizes))>::type deallocateImplem (void* p) noexcept
        {
            auto& pool = std::get<I> (_pools);
            if (pool.owns (p))
            {
                pool.push (p);
                return;
            }
            deallocateImplem<I + 1> (p);
        }
 
        template <size_t I>
        typename std::enable_if<(I >= sizeof...(Sizes))>::type deallocateImplem (void*) noexcept
        {
        }
 
        Backend _backend;
 
        std::tuple<BasicPool<Count, Sizes>...> _pools;
    };
}
 
#endif