numa_work_stealer.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 "enhanced_work_stealing_config.h"
#include "numa_topology.h"
#include "steal_backoff_strategy.h"
#include "work_affinity_tracker.h"
#include "work_stealing_stats.h"
 
#include <cstddef>
#include <functional>
#include <memory>
#include <random>
#include <vector>
 
namespace kcenon::thread
{
 
// Forward declarations
class thread_worker;
class job;
 
namespace lockfree
{
template <typename T>
class work_stealing_deque;
} // namespace lockfree
 
class numa_work_stealer
{
public:
    using deque_accessor_fn = std::function<lockfree::work_stealing_deque<job*>*(std::size_t)>;
 
    using cpu_accessor_fn = std::function<int(std::size_t)>;
 
    numa_work_stealer(std::size_t worker_count,
                      deque_accessor_fn deque_accessor,
                      cpu_accessor_fn cpu_accessor,
                      enhanced_work_stealing_config config = {});
 
    ~numa_work_stealer() = default;
 
    // Non-copyable, non-movable (contains atomics)
    numa_work_stealer(const numa_work_stealer&) = delete;
    numa_work_stealer& operator=(const numa_work_stealer&) = delete;
    numa_work_stealer(numa_work_stealer&&) = delete;
    numa_work_stealer& operator=(numa_work_stealer&&) = delete;
 
    [[nodiscard]] auto steal_for(std::size_t worker_id) -> job*;
 
    [[nodiscard]] auto steal_batch_for(std::size_t worker_id, std::size_t max_count)
        -> std::vector<job*>;
 
    [[nodiscard]] auto get_stats() const -> const work_stealing_stats&;
 
    [[nodiscard]] auto get_stats_snapshot() const -> work_stealing_stats_snapshot;
 
    void reset_stats();
 
    [[nodiscard]] auto get_topology() const -> const numa_topology&;
 
    [[nodiscard]] auto get_config() const -> const enhanced_work_stealing_config&;
 
    void set_config(const enhanced_work_stealing_config& config);
 
private:
    [[nodiscard]] auto select_victims(std::size_t requester_id, std::size_t count)
        -> std::vector<std::size_t>;
 
    [[nodiscard]] auto select_victims_random(std::size_t requester_id, std::size_t count)
        -> std::vector<std::size_t>;
 
    [[nodiscard]] auto select_victims_round_robin(std::size_t requester_id, std::size_t count)
        -> std::vector<std::size_t>;
 
    [[nodiscard]] auto select_victims_adaptive(std::size_t requester_id, std::size_t count)
        -> std::vector<std::size_t>;
 
    [[nodiscard]] auto select_victims_numa_aware(std::size_t requester_id, std::size_t count)
        -> std::vector<std::size_t>;
 
    [[nodiscard]] auto select_victims_locality_aware(std::size_t requester_id, std::size_t count)
        -> std::vector<std::size_t>;
 
    [[nodiscard]] auto select_victims_hierarchical(std::size_t requester_id, std::size_t count)
        -> std::vector<std::size_t>;
 
    [[nodiscard]] auto calculate_batch_size(std::size_t victim_queue_size) const -> std::size_t;
 
    [[nodiscard]] auto get_worker_cpu(std::size_t worker_id) const -> int;
 
    [[nodiscard]] auto workers_on_same_node(std::size_t worker_a, std::size_t worker_b) const
        -> bool;
 
    void record_steal(std::size_t thief_id, std::size_t victim_id);
 
    std::size_t worker_count_;
    deque_accessor_fn deque_accessor_;
    cpu_accessor_fn cpu_accessor_;
    enhanced_work_stealing_config config_;
    numa_topology topology_;
    work_stealing_stats stats_;
    std::unique_ptr<work_affinity_tracker> affinity_tracker_;
    std::unique_ptr<backoff_calculator> backoff_calculator_;
 
    // Per-thread random generators for random victim selection
    mutable std::mt19937_64 rng_;
 
    // Round-robin state
    mutable std::atomic<std::size_t> round_robin_index_{0};
};
 
} // namespace kcenon::thread