thread_pool.hpp

Go to the documentation of this file.

#pragma once

 
#include <condition_variable>
#include <cstdint>
#include <mutex>
#include <thread>
#include <vector>
 
#include "orrery/backend/cpu_topology.hpp"
#include "orrery/core/function_ref.hpp"
 
namespace orrery::backend {

class ThreadPool {
public:
    enum class Affinity : std::uint8_t {
        kUnpinned,
        kPinned,
    };

    using WorkerBody = core::FunctionRef<void(unsigned)>;

    explicit ThreadPool(unsigned worker_count = 0, Affinity affinity = Affinity::kUnpinned);

    ~ThreadPool();
 
    // A pool owns running threads that hold a pointer to it, so it cannot be
    // copied and cannot move out from under them.
    ThreadPool(const ThreadPool&) = delete;
    ThreadPool& operator=(const ThreadPool&) = delete;
    ThreadPool(ThreadPool&&) = delete;
    ThreadPool& operator=(ThreadPool&&) = delete;

    void dispatch(WorkerBody body);
 
    [[nodiscard]] unsigned worker_count() const noexcept {
        return static_cast<unsigned>(threads_.size());
    }
 
    [[nodiscard]] Affinity affinity() const noexcept { return affinity_; }

    [[nodiscard]] CoreClass core_class(unsigned worker) const noexcept;

    [[nodiscard]] static unsigned default_worker_count() noexcept;
 
private:
    void worker_loop(unsigned worker) noexcept;
 
    std::mutex mutex_;

    std::condition_variable work_available_;

    std::condition_variable work_complete_;

    std::condition_variable worker_started_;

    const WorkerBody* body_{nullptr};

    std::uint64_t generation_{0};

    unsigned pending_{0};

    unsigned starting_{0};
 
    bool stopping_{false};
 
    Affinity affinity_{Affinity::kUnpinned};

    std::vector<unsigned> pin_targets_;

    std::vector<CoreClass> core_classes_;

    std::vector<std::thread> threads_;
};
 
} // namespace orrery::backend