AtomicPool.hpp

Go to the documentation of this file.

#pragma once
#include <cstring>
#include "Allocator.hpp"
#include "Atomic.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;
            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));
            node->used = false;
            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(); }
    };
} // namespace Foundation::Core