AtomicPool.hpp

Go to the documentation of this file.

#pragma once
#include <cstddef>
#include <cstring>
#include <memory>
#include <mutex>
#include <utility>
#include "Allocator.hpp"
#include "Atomic.hpp"
#include "Container.hpp"
#include "Thread.hpp"
namespace Foundation::Core
{
    template <typename T>        
    class AtomicPool
    {
        struct Node;
        struct alignas(2 * sizeof(Node*)) PTag
        {
            Node* p{nullptr};
            uintptr_t tag{};
        };
        struct alignas(2 * sizeof(Node*)) Node
        {
            Node* next;
            Atomic<uintptr_t> used{0};
            T data;            
            
        };
        Node* mNodes{nullptr};
        size_t mSize{0};
        Atomic<PTag> mHead{};
        ScopedArena mArena;
 
        Node* AllocateNode()
        {
            PTag old_head = mHead.load(std::memory_order_relaxed);
            while (true)
            {
                if (!old_head.p)
                    return nullptr;
                // ABA mitigation as seen in AtomicStack
                PTag new_head = {old_head.p->next, old_head.tag + 1};
                if (mHead.compare_exchange_weak(old_head, new_head, std::memory_order_acquire,
                                                std::memory_order_relaxed))
                {
                    return old_head.p;
                }
            }
        }
 
        void DeallocateNode(Node* node)
        {
            PTag old_head = mHead.load(std::memory_order_relaxed);
            while (true)
            {
                node->next = old_head.p;
                PTag new_head{node, old_head.tag + 1};
                if (mHead.compare_exchange_weak(old_head, new_head, std::memory_order_acquire,
                                                std::memory_order_relaxed))
                {
                    return;
                }
            }
        }
 
    public:
        AtomicPool() = default;
        AtomicPool(size_t size, Allocator* alloc) : mSize(size), mArena(alloc, size * sizeof(Node))
        {
            std::memset(mArena.arena.memory, 0, mArena.arena.size);
            mNodes = static_cast<Node*>(mArena.arena.memory);
            for (size_t i = 0; i + 1 < mSize; i++)
                mNodes[i].next = &mNodes[i + 1];
            mNodes[mSize - 1].next = nullptr;
            mHead.store({&mNodes[0], 0});
        }
        template <typename... Args>
        T* Construct(Args&&... args)
        {
            Node* node = AllocateNode();
            if (!node)
                return nullptr;
            node->used = true;
            std::construct_at(&node->data, std::forward<Args>(args)...);
            return &node->data;
        }
        void Destruct(T* ptr)
        {
            if (!ptr)
                return;
            // We know the memory layout is Node.data, so we can get the Node* from T*
            // Hacky - though guaranteed to work by standard layout.
            Node* node = reinterpret_cast<Node*>(reinterpret_cast<uintptr_t>(ptr) - offsetof(Node, data));
            uintptr_t old_used = true;
            if (node->used.compare_exchange_strong(old_used, false, std::memory_order_acquire, std::memory_order_relaxed))
            {
                node->data.~T();
                DeallocateNode(node);
            }
        }
        size_t Index(T* ptr) const
        {
            Node* node = reinterpret_cast<Node*>(reinterpret_cast<uintptr_t>(ptr) - offsetof(Node, data));
            return static_cast<size_t>(node - mNodes);
        }
        T* At(size_t index)
        {
            if (index >= mSize)
                return nullptr;
            return &mNodes[index].data;
        }
        void Collect()
        {
            for (size_t i = 0; i < mSize; i++)
                if (mNodes[i].used)
                    mNodes[i].used = false, mNodes[i].data.~T();
        }
        ~AtomicPool() { Collect(); }
    };
 
    template <typename T>
    class DynamicPool
    {
        Allocator* mAllocator;
        List<T> mObjects;
        mutable Mutex mMutex;
 
    public:
        explicit DynamicPool(Allocator* alloc) : mAllocator(alloc), mObjects(alloc) {}
 
        DynamicPool(DynamicPool const&) = delete;
        DynamicPool& operator=(DynamicPool const&) = delete;
        DynamicPool(DynamicPool&&) = delete;
        DynamicPool& operator=(DynamicPool&&) = delete;
 
        template <typename... Args>
        T* Construct(Args&&... args)
        {
            std::unique_lock lock(mMutex);
            return &mObjects.emplace_back(std::forward<Args>(args)...);
        }
 
        void Destruct(T* ptr)
        {
            if (!ptr)
                return;
 
            List<T> removed(mAllocator);
            {
                std::unique_lock lock(mMutex);
                for (auto it = mObjects.begin(); it != mObjects.end(); ++it)
                {
                    if (&*it == ptr)
                    {
                        removed.splice(removed.end(), mObjects, it);
                        break;
                    }
                }
            }          
        }
 
        void Collect()
        {
            List<T> removed(mAllocator);
            {
                std::unique_lock lock(mMutex);
                removed.splice(removed.end(), mObjects);
            }            
        }
 
        [[nodiscard]] size_t Size() const
        {
            std::unique_lock lock(mMutex);
            return mObjects.size();
        }
 
        ~DynamicPool() { Collect(); }
    };
} // namespace Foundation::Core