work stealing queue

workstealingqueue.cpp

#include <atomic>

class Job;

struct WorkStealingQueue
{
    // only called by owner work thread
    void Push(Job* job) noexcept
    {
        // m_bottom -> stealing thread read, owner thread read, write.
        auto bottom = m_bottom.load(std::memory_order_relaxed);

        m_jobs[bottom & MASK] = job;

        // need  let stealing thread see the new job.
        m_bottom.store(bottom + 1, std::memory_order_release);
    }

    // only called by owner worker thread 
    Job* Pop(void) noexcept
    {
        auto bottom = m_bottom.fetch_sub(1, std::memory_order_relaxed);

        auto top = m_top.load(std::memory_order_acquire);

        if (bottom >= top)
        {
            auto job = m_jobs[bottom  & MASK];

            if(top != bottom) 
            {
                return job;
            }
            auto top1 = top + 1;
            if (!m_top.compare_exchange_weak(
                top, top1,
                std::memory_order_release,
                std::memory_order_relaxed))
            {
                job = nullptr;
            } 
            m_bottom.store(top1);
            return job;
        }
        m_bottom.store(top, std::memory_order_release);
        return nullptr;
    }

    // called by stealing thread, not owner thread
    Job* Steal(void) noexcept
    {
        auto top = m_top.load(std::memory_order_acquire);

        // Release-Acquire ordering, so the stealing thread see new job
        auto bottom = m_bottom.load(std::memory_order_acquire);

        if (bottom > top)
        {
            auto job = m_jobs[top & MASK];
            // check if other stealing thread stealing this work 
            // or owner thread pop this job.
            // no data should do sync, so use relaxed oreder
            if (m_top.compare_exchange_weak(
                top, top + 1,
                std::memory_order_release,
                std::memory_order_relaxed))
            {
                return job;
            }
        }

        return nullptr;
    }

private:

    static constexpr auto MAX_COUNT = 4096u;
    static constexpr auto MASK = MAX_COUNT - 1u;
    static_assert((MAX_COUNT & MASK) == 0, "the max number of job must be power of two.");

    Job* m_jobs[MAX_COUNT];
    std::atomic<unsigned int> m_bottom{ 0 }, m_top{ 0 };
};