work_stealing_deque.h

Go to the documentation of this file.

// BSD 3-Clause License
// Copyright (c) 2024, 🍀☀🌕🌥 🌊
// See the LICENSE file in the project root for full license information.
 
#pragma once
 
#include <algorithm>
#include <atomic>
#include <cstdint>
#include <memory>
#include <optional>
#include <vector>
 
namespace kcenon::thread::lockfree {
 
template<typename T>
class circular_array {
public:
    explicit circular_array(std::size_t log_size)
        : log_size_(log_size)
        , size_(1ULL << log_size)
        , mask_(size_ - 1)
        , buffer_(new std::atomic<T>[size_]) {
        for (std::size_t i = 0; i < size_; ++i) {
            buffer_[i].store(T{}, std::memory_order_relaxed);
        }
    }
 
    ~circular_array() {
        delete[] buffer_;
    }
 
    // Non-copyable
    circular_array(const circular_array&) = delete;
    circular_array& operator=(const circular_array&) = delete;
 
    [[nodiscard]] std::size_t size() const noexcept {
        return size_;
    }
 
    [[nodiscard]] T get(std::int64_t index) const noexcept {
        return buffer_[index & mask_].load(std::memory_order_relaxed);
    }
 
    void put(std::int64_t index, T value) noexcept {
        buffer_[index & mask_].store(value, std::memory_order_relaxed);
    }
 
    [[nodiscard]] circular_array* grow(std::int64_t bottom, std::int64_t top) const {
        auto* new_array = new circular_array(log_size_ + 1);
        for (std::int64_t i = top; i < bottom; ++i) {
            new_array->put(i, get(i));
        }
        return new_array;
    }
 
private:
    std::size_t log_size_;
    std::size_t size_;
    std::size_t mask_;
    std::atomic<T>* buffer_;
};
 
template<typename T>
class work_stealing_deque {
public:
    static constexpr std::size_t LOG_INITIAL_SIZE = 5;
 
    explicit work_stealing_deque(std::size_t log_initial_size = LOG_INITIAL_SIZE)
        : top_(0)
        , bottom_(0)
        , array_(new circular_array<T>(log_initial_size)) {
    }
 
    ~work_stealing_deque() {
        delete array_.load(std::memory_order_relaxed);
    }
 
    // Non-copyable and non-movable
    work_stealing_deque(const work_stealing_deque&) = delete;
    work_stealing_deque& operator=(const work_stealing_deque&) = delete;
    work_stealing_deque(work_stealing_deque&&) = delete;
    work_stealing_deque& operator=(work_stealing_deque&&) = delete;
 
    void push(T item) {
        std::int64_t b = bottom_.load(std::memory_order_relaxed);
        std::int64_t t = top_.load(std::memory_order_acquire);
        circular_array<T>* a = array_.load(std::memory_order_relaxed);
 
        // Check if array needs to grow
        if (b - t > static_cast<std::int64_t>(a->size()) - 1) {
            // Grow the array
            circular_array<T>* new_array = a->grow(b, t);
            // Store old array for cleanup (in a real implementation,
            // you would use hazard pointers or epoch-based reclamation)
            old_arrays_.push_back(a);
            array_.store(new_array, std::memory_order_release);
            a = new_array;
        }
 
        a->put(b, item);
        std::atomic_thread_fence(std::memory_order_release);
        bottom_.store(b + 1, std::memory_order_relaxed);
    }
 
    [[nodiscard]] std::optional<T> pop() {
        std::int64_t b = bottom_.load(std::memory_order_relaxed) - 1;
        circular_array<T>* a = array_.load(std::memory_order_relaxed);
        bottom_.store(b, std::memory_order_relaxed);
        std::atomic_thread_fence(std::memory_order_seq_cst);
        std::int64_t t = top_.load(std::memory_order_relaxed);
 
        if (t <= b) {
            // Non-empty queue
            T item = a->get(b);
            if (t == b) {
                // Last element - compete with thieves
                if (!top_.compare_exchange_strong(
                        t, t + 1,
                        std::memory_order_seq_cst,
                        std::memory_order_relaxed)) {
                    // Lost race with a thief
                    bottom_.store(b + 1, std::memory_order_relaxed);
                    return std::nullopt;
                }
                bottom_.store(b + 1, std::memory_order_relaxed);
            }
            return item;
        } else {
            // Empty queue
            bottom_.store(b + 1, std::memory_order_relaxed);
            return std::nullopt;
        }
    }
 
    [[nodiscard]] std::optional<T> steal() {
        std::int64_t t = top_.load(std::memory_order_acquire);
        std::atomic_thread_fence(std::memory_order_seq_cst);
        std::int64_t b = bottom_.load(std::memory_order_acquire);
 
        if (t < b) {
            // Non-empty queue
            circular_array<T>* a = array_.load(std::memory_order_consume);
            T item = a->get(t);
 
            if (!top_.compare_exchange_strong(
                    t, t + 1,
                    std::memory_order_seq_cst,
                    std::memory_order_relaxed)) {
                // Lost race with another thief or owner
                return std::nullopt;
            }
            return item;
        }
        return std::nullopt;
    }
 
    [[nodiscard]] std::vector<T> steal_batch(std::size_t max_count) {
        if (max_count == 0) {
            return {};
        }
 
        std::int64_t t = top_.load(std::memory_order_acquire);
        std::atomic_thread_fence(std::memory_order_seq_cst);
        std::int64_t b = bottom_.load(std::memory_order_acquire);
 
        if (t >= b) {
            // Empty queue
            return {};
        }
 
        // Calculate how many we can actually steal
        std::int64_t available = b - t;
        std::size_t to_steal = std::min(
            max_count,
            static_cast<std::size_t>(available)
        );
 
        // Try to atomically claim the range [t, t + to_steal)
        std::int64_t new_top = t + static_cast<std::int64_t>(to_steal);
 
        if (!top_.compare_exchange_strong(
                t, new_top,
                std::memory_order_seq_cst,
                std::memory_order_relaxed)) {
            // Lost race with another thief or owner
            // Return empty and let caller retry if needed
            return {};
        }
 
        // Successfully claimed the range - now read the elements
        // The CAS already ensured we have exclusive access to [t, new_top)
        circular_array<T>* a = array_.load(std::memory_order_consume);
        std::vector<T> result;
        result.reserve(to_steal);
 
        for (std::int64_t i = t; i < new_top; ++i) {
            result.push_back(a->get(i));
        }
 
        return result;
    }
 
    [[nodiscard]] bool empty() const noexcept {
        std::int64_t b = bottom_.load(std::memory_order_relaxed);
        std::int64_t t = top_.load(std::memory_order_relaxed);
        return b <= t;
    }
 
    [[nodiscard]] std::size_t size() const noexcept {
        std::int64_t b = bottom_.load(std::memory_order_relaxed);
        std::int64_t t = top_.load(std::memory_order_relaxed);
        std::int64_t diff = b - t;
        return diff > 0 ? static_cast<std::size_t>(diff) : 0;
    }
 
    [[nodiscard]] std::size_t capacity() const noexcept {
        return array_.load(std::memory_order_relaxed)->size();
    }
 
    void cleanup_old_arrays() {
        for (auto* old_array : old_arrays_) {
            delete old_array;
        }
        old_arrays_.clear();
    }
 
private:
    // Cache line padding to prevent false sharing
    static constexpr std::size_t CACHE_LINE_SIZE = 64;
 
    alignas(CACHE_LINE_SIZE) std::atomic<std::int64_t> top_;
    alignas(CACHE_LINE_SIZE) std::atomic<std::int64_t> bottom_;
    alignas(CACHE_LINE_SIZE) std::atomic<circular_array<T>*> array_;
 
    // Storage for old arrays (simple approach - could use hazard pointers)
    std::vector<circular_array<T>*> old_arrays_;
};
 
} // namespace kcenon::thread::lockfree