SlabAllocator.h

// Copyright (c) 2022, Nikolaos D. Bougalis <nikb@bougalis.net>
 
#pragma once
 
#include <xrpl/basics/ByteUtilities.h>
#include <xrpl/beast/type_name.h>
#include <xrpl/beast/utility/instrumentation.h>
 
#include <boost/align.hpp>
#include <boost/container/static_vector.hpp>
#include <boost/predef.h>
 
#include <algorithm>
#include <atomic>
#include <cstdint>
#include <cstring>
#include <mutex>
#include <vector>
 
#if BOOST_OS_LINUX
#include <sys/mman.h>
#endif
 
namespace xrpl {
 
template <typename Type>
class SlabAllocator
{
    static_assert(
        sizeof(Type) >= sizeof(std::uint8_t*),
        "SlabAllocator: the requested object must be larger than a pointer.");
 
    static_assert(alignof(Type) == 8 || alignof(Type) == 4);

    struct SlabBlock
    {
        // A mutex to protect the freelist for this block:
        std::mutex m;
 
        // A linked list of appropriately sized free buffers:
        std::uint8_t* l = nullptr;
 
        // The next memory block
        SlabBlock* next;
 
        // The underlying memory block:
        std::uint8_t const* const p = nullptr;
 
        // The extent of the underlying memory block:
        std::size_t const size;
 
        SlabBlock(SlabBlock* next, std::uint8_t* data, std::size_t size, std::size_t item)
            : next(next), p(data), size(size)
        {
            // We don't need to grab the mutex here, since we're the only
            // ones with access at this moment.
 
            while (data + item <= p + size)
            {
                // Use memcpy to avoid unaligned UB
                // (will optimize to equivalent code)
                std::memcpy(data, static_cast<void const*>(&l), sizeof(std::uint8_t*));
                l = data;
                data += item;
            }
        }
 
        // Calling this destructor will release the allocated memory but
        // will not properly destroy any objects that are constructed in
        // the block itself.
        ~SlabBlock() = default;
 
        SlabBlock(SlabBlock const& other) = delete;
        SlabBlock&
        operator=(SlabBlock const& other) = delete;
 
        SlabBlock(SlabBlock&& other) = delete;
        SlabBlock&
        operator=(SlabBlock&& other) = delete;

        bool
        own(std::uint8_t const* pIn) const noexcept
        {
            return (pIn >= p) && (pIn < p + size);
        }
 
        std::uint8_t*
        allocate() noexcept
        {
            std::uint8_t* ret = nullptr;  // NOLINT(misc-const-correctness)
 
            {
                std::scoped_lock const lock(m);
                ret = l;
 
                if (ret != nullptr)
                {
                    // Use memcpy to avoid unaligned UB
                    // (will optimize to equivalent code)
                    std::memcpy(static_cast<void*>(&l), ret, sizeof(std::uint8_t*));
                }
            }
 
            return ret;
        }

        void
        deallocate(std::uint8_t* ptr) noexcept
        {
            XRPL_ASSERT(own(ptr), "xrpl::SlabAllocator::SlabBlock::deallocate : own input");
 
            std::scoped_lock const lock(m);
 
            // Use memcpy to avoid unaligned UB
            // (will optimize to equivalent code)
            std::memcpy(ptr, static_cast<void const*>(&l), sizeof(std::uint8_t*));
            l = ptr;
        }
    };
 
private:
    // A linked list of slabs
    std::atomic<SlabBlock*> slabs_ = nullptr;
 
    // The alignment requirements of the item we're allocating:
    std::size_t const itemAlignment_;
 
    // The size of an item, including the extra bytes requested and
    // any padding needed for alignment purposes:
    std::size_t const itemSize_;
 
    // The size of each individual slab:
    std::size_t const slabSize_;
 
public:
    constexpr explicit SlabAllocator(
        std::size_t extra,
        std::size_t alloc = 0,
        std::size_t align = 0)
        : itemAlignment_((align != 0u) ? align : alignof(Type))
        , itemSize_(boost::alignment::align_up(sizeof(Type) + extra, itemAlignment_))
        , slabSize_(alloc)
    {
        XRPL_ASSERT(
            (itemAlignment_ & (itemAlignment_ - 1)) == 0,
            "xrpl::SlabAllocator::SlabAllocator : valid alignment");
    }
 
    SlabAllocator(SlabAllocator const& other) = delete;
    SlabAllocator&
    operator=(SlabAllocator const& other) = delete;
 
    SlabAllocator(SlabAllocator&& other) = delete;
    SlabAllocator&
    operator=(SlabAllocator&& other) = delete;
 
    // FIXME: We can't destroy the memory blocks we've allocated, because
    //        we can't be sure that they are not being used. Cleaning the
    //        shutdown process up could make this possible.
    ~SlabAllocator() = default;

    [[nodiscard]] constexpr std::size_t
    size() const noexcept
    {
        return itemSize_;
    }

    std::uint8_t*
    allocate() noexcept
    {
        auto slab = slabs_.load();
 
        while (slab != nullptr)
        {
            if (auto ret = slab->allocate())
                return ret;
 
            slab = slab->next;
        }
 
        // No slab can satisfy our request, so we attempt to allocate a new
        // one here:
        std::size_t const size = slabSize_;
 
        // We want to allocate the memory at a 2 MiB boundary, to make it
        // possible to use hugepage mappings on Linux:
        auto buf = boost::alignment::aligned_alloc(megabytes(std::size_t(2)), size);
        if (buf == nullptr) [[unlikely]]
            return nullptr;
 
#if BOOST_OS_LINUX
        // When allocating large blocks, attempt to leverage Linux's
        // transparent hugepage support. It is unclear and difficult
        // to accurately determine if doing this impacts performance
        // enough to justify using platform-specific tricks.
        if (size >= megabytes(std::size_t(4)))
            madvise(buf, size, MADV_HUGEPAGE);
#endif
 
        // We need to carve out a bit of memory for the slab header
        // and then align the rest appropriately:
        auto slabData =
            reinterpret_cast<void*>(reinterpret_cast<std::uint8_t*>(buf) + sizeof(SlabBlock));
        auto slabSize = size - sizeof(SlabBlock);
 
        // This operation is essentially guaranteed not to fail but
        // let's be careful anyways.
        if (boost::alignment::align(itemAlignment_, itemSize_, slabData, slabSize) == nullptr)
        {
            boost::alignment::aligned_free(buf);
            return nullptr;
        }
 
        slab = new (buf) SlabBlock(
            slabs_.load(), reinterpret_cast<std::uint8_t*>(slabData), slabSize, itemSize_);
 
        // Link the new slab
        while (!slabs_.compare_exchange_weak(
            slab->next, slab, std::memory_order_release, std::memory_order_relaxed))
        {
            ;  // Nothing to do
        }
 
        return slab->allocate();
    }

    bool
    deallocate(std::uint8_t* ptr) noexcept
    {
        XRPL_ASSERT(
            ptr,
            "xrpl::SlabAllocator::SlabAllocator::deallocate : non-null "
            "input");
 
        for (auto slab = slabs_.load(); slab != nullptr; slab = slab->next)
        {
            if (slab->own(ptr))
            {
                slab->deallocate(ptr);
                return true;
            }
        }
 
        return false;
    }
};

template <typename Type>
class SlabAllocatorSet
{
private:
    // The list of allocators that belong to this set
    boost::container::static_vector<SlabAllocator<Type>, 64> allocators_{};
 
    std::size_t maxSize_ = 0;
 
public:
    class SlabConfig
    {
        friend class SlabAllocatorSet;
 
    private:
        std::size_t extra_;
        std::size_t alloc_;
        std::size_t align_;
 
    public:
        constexpr SlabConfig(
            std::size_t extra,
            std::size_t alloc = 0,
            std::size_t align = alignof(Type))
            : extra_(extra), alloc_(alloc), align_(align)
        {
        }
    };
 
    constexpr SlabAllocatorSet(std::vector<SlabConfig> cfg)
    {
        // Ensure that the specified allocators are sorted from smallest to
        // largest by size:
        std::sort(std::begin(cfg), std::end(cfg), [](SlabConfig const& a, SlabConfig const& b) {
            return a.extra_ < b.extra_;
        });
 
        // We should never have two slabs of the same size
        if (std::adjacent_find(
                std::begin(cfg), std::end(cfg), [](SlabConfig const& a, SlabConfig const& b) {
                    return a.extra_ == b.extra_;
                }) != cfg.end())
        {
            throw std::runtime_error(
                "SlabAllocatorSet<" + beast::typeName<Type>() + ">: duplicate slab size");
        }
 
        for (auto const& c : cfg)
        {
            auto& a = allocators_.emplace_back(c.extra_, c.alloc_, c.align_);
 
            if (a.size() > maxSize_)
                maxSize_ = a.size();
        }
    }
 
    SlabAllocatorSet(SlabAllocatorSet const& other) = delete;
    SlabAllocatorSet&
    operator=(SlabAllocatorSet const& other) = delete;
 
    SlabAllocatorSet(SlabAllocatorSet&& other) = delete;
    SlabAllocatorSet&
    operator=(SlabAllocatorSet&& other) = delete;
 
    ~SlabAllocatorSet() = default;

    std::uint8_t*
    allocate(std::size_t extra) noexcept
    {
        if (auto const size = sizeof(Type) + extra; size <= maxSize_)
        {
            for (auto& a : allocators_)
            {
                if (a.size() >= size)
                    return a.allocate();
            }
        }
 
        return nullptr;
    }

    bool
    deallocate(std::uint8_t* ptr) noexcept
    {
        for (auto& a : allocators_)
        {
            if (a.deallocate(ptr))
                return true;
        }
 
        return false;
    }
};
 
}  // namespace xrpl