AtomicQueue.hpp

Go to the documentation of this file.

#pragma once
#include <bit>
#include "Allocator.hpp"
#include "Atomic.hpp"
#include "Container.hpp"
#include "Logging.hpp"
namespace Foundation::Core
{
    template <typename T>
    class SPSCQueue
    {
        const size_t mModulo;
        Vector<T> mBuffer;
        Atomic<size_t> mRead{}, mWrite{};
        // Only used in reader thread
        size_t mReadCached{};
        // Only used in writer thread
        size_t mWriteCached{};
 
    public:
        SPSCQueue(size_t size, Allocator* alloc) : mModulo(size - 1), mBuffer(size, alloc)
        {
            CHECK_MSG((size & mModulo) == 0, "Size must be a power of two");
        }
        template <typename U>
        bool Push(U&& data)
        {
            size_t write = mWrite.load(std::memory_order_relaxed);
            // Amortize atomic reads
            // We don't need to update the read head everytime. Only when we think we're full.
            if (write - mReadCached == mBuffer.size()) [[unlikely]]
            {
                // Wrapped. Next cycle could begin only when there's enough space
                mReadCached = mRead.load(std::memory_order_acquire);
                if (write - mReadCached == mBuffer.size())
                    return false; // full
            }
            mBuffer[write & mModulo] = std::forward<U>(data);
            mWrite.store(write + 1, std::memory_order_release);
            return true;
        }
        bool Pop(T& out)
        {
            size_t read = mRead.load(std::memory_order_relaxed);
            // Same as above
            if (read == mWriteCached) [[unlikely]]
            {
                mWriteCached = mWrite.load(std::memory_order_acquire);
                if (read == mWriteCached)
                    return false; // empty
            }
            mRead.store(read + 1, std::memory_order_release);
            out = std::move(mBuffer[read & mModulo]);
            return true;
        }
    };
    template <typename T>
    class MPMCQueue
    {
        struct Data
        {
            T data{};
            Atomic<size_t> writeCycle{} /* when to write */, readCycle{} /* when to read */;
        };
        const size_t mModulo, mShift;
        Vector<Data> mBuffer;
        Atomic<size_t> mRead{}, mWrite{};
        size_t mWriteCached{};
 
    public:
        MPMCQueue(size_t size, Allocator* alloc) :
            mModulo(size - 1), mShift(std::countr_zero(size)), mBuffer(size, alloc)
        {
            CHECK_MSG((size & mModulo) == 0, "Size must be a power of two");
        }
        template <typename U>
        bool Push(U&& data)
        {
            size_t write = mWrite.load(std::memory_order_relaxed);
            while (true)
            {
                auto& elem = mBuffer[write & mModulo];
                size_t cycle = write >> mShift;
                size_t read_cycle = elem.readCycle.load(std::memory_order_acquire);
                // Waits for consumer to pop, if any
                if (read_cycle == cycle) [[likely]] // Ready to write
                {
                    // Bump the write index if we can, claiming the old index. Try later otherwise.
                    if (mWrite.compare_exchange_weak(write, write + 1, std::memory_order_relaxed))
                    {
                        elem.data = std::forward<U>(data);
                        elem.writeCycle.store(cycle + 1, std::memory_order_release);
                        return true;
                    }
                }
                else if (read_cycle >= cycle) // Not our turn yet? So we must be an old write. Update and try again. CAS
                                              // does this already.
                {
                    write = mWrite.load(std::memory_order_relaxed);
                }
                else // Counter is behind write cycle! We're full
                {
                    return false;
                }
            }
        }
        bool Pop(T& out)
        {
            size_t read = mRead.load(std::memory_order_relaxed);
            while (true)
            {
                // Same as above, but wait for producer
                auto& elem = mBuffer[read & mModulo];
                size_t cycle = read >> mShift;
                size_t write_cycle = elem.writeCycle.load(std::memory_order_acquire);
                if (write_cycle == cycle + 1)
                {
                    // Same as above for reading
                    if (mRead.compare_exchange_weak(read, read + 1, std::memory_order_relaxed))
                    {
                        out = std::move(elem.data);
                        elem.readCycle.store(cycle + 1, std::memory_order_release);
                        return true;
                    }
                }
                else if (write_cycle > cycle) // Old read, not synced so try again.
                {
                    read = mRead.load(std::memory_order_relaxed);
                }
                else // Read cycle is behind write cycle! We're empty
                {
                    return false;
                }
            }
        }
    };
} // namespace Foundation::Core