212 lines
8.2 KiB
C++
212 lines
8.2 KiB
C++
#pragma once
|
|
#include "diagnostics.hpp"
|
|
#include <atomic>
|
|
#include <condition_variable>
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <optional>
|
|
#include <queue>
|
|
#include <thread>
|
|
#include <vector>
|
|
|
|
namespace kpn {
|
|
|
|
// ── IScheduler ────────────────────────────────────────────────────────────────
|
|
|
|
struct IScheduler {
|
|
virtual ~IScheduler() = default;
|
|
|
|
// Submit a task with an optional priority in [0, 1]. Higher = run sooner.
|
|
virtual void submit(std::function<void()> task, float priority = 0.5f) = 0;
|
|
|
|
// Start worker threads. Must be called before submit().
|
|
virtual void start() = 0;
|
|
|
|
// Halt: signal workers to exit and join them. Pending tasks are discarded.
|
|
virtual void stop() = 0;
|
|
|
|
// Drain: block until all in-flight tasks complete. Workers keep running.
|
|
virtual void drain() = 0;
|
|
};
|
|
|
|
// ── ThreadPool ────────────────────────────────────────────────────────────────
|
|
//
|
|
// Work-stealing thread pool with per-thread priority queues.
|
|
//
|
|
// Each worker owns a priority_queue (max-heap by priority, FIFO within equal
|
|
// priority via sequence number). submit() distributes via round-robin. When a
|
|
// worker's queue is empty it tries to steal from the most-loaded peer using
|
|
// try_lock to avoid blocking; if no work is found it sleeps on a shared CV.
|
|
//
|
|
// total_ counts tasks submitted-but-not-completed (queued + executing).
|
|
// drain() waits until total_ == 0.
|
|
|
|
class ThreadPool : public IScheduler, public IPoolProbe {
|
|
public:
|
|
explicit ThreadPool(std::size_t thread_count) : thread_count_(thread_count) {}
|
|
|
|
~ThreadPool() {
|
|
if (!stopped_.load(std::memory_order_relaxed))
|
|
stop();
|
|
}
|
|
|
|
void start() override {
|
|
stopped_.store(false, std::memory_order_relaxed);
|
|
queues_.clear();
|
|
for (std::size_t i = 0; i < thread_count_; ++i)
|
|
queues_.push_back(std::make_unique<WorkerQueue>());
|
|
workers_.reserve(thread_count_);
|
|
for (std::size_t i = 0; i < thread_count_; ++i)
|
|
workers_.emplace_back([this, i] { worker_loop(i); });
|
|
}
|
|
|
|
void stop() override {
|
|
stopped_.store(true, std::memory_order_seq_cst);
|
|
for (auto& q : queues_) {
|
|
std::lock_guard lock(q->mx);
|
|
std::size_t discarded = q->pq.size();
|
|
while (!q->pq.empty()) q->pq.pop();
|
|
total_.fetch_sub(discarded, std::memory_order_relaxed);
|
|
}
|
|
// Lock cv_mx_ before notifying so the stop signal can't be lost in the
|
|
// gap between a worker's predicate check and its wait() (see submit()).
|
|
{ std::lock_guard<std::mutex> lk(cv_mx_); }
|
|
cv_.notify_all();
|
|
for (auto& t : workers_) if (t.joinable()) t.join();
|
|
workers_.clear();
|
|
queues_.clear();
|
|
}
|
|
|
|
void drain() override {
|
|
std::unique_lock lock(drain_mx_);
|
|
drain_cv_.wait(lock, [this] {
|
|
return total_.load(std::memory_order_acquire) == 0;
|
|
});
|
|
}
|
|
|
|
void submit(std::function<void()> task, float priority = 0.5f) override {
|
|
std::size_t target = next_.fetch_add(1, std::memory_order_relaxed) % thread_count_;
|
|
{
|
|
std::lock_guard lock(queues_[target]->mx);
|
|
queues_[target]->pq.push(
|
|
{std::move(task), priority, seq_.fetch_add(1, std::memory_order_relaxed)});
|
|
}
|
|
total_.fetch_add(1, std::memory_order_relaxed);
|
|
submitted_.fetch_add(1, std::memory_order_relaxed);
|
|
// Synchronize with worker_loop's predicate evaluation: taking cv_mx_
|
|
// here guarantees a worker is either before its predicate check (and
|
|
// will observe total_ > 0) or already blocked in wait() (and will be
|
|
// woken). Without this, notify_one() can slip into the gap between the
|
|
// worker's predicate check and its wait(), and be lost — a deadlock.
|
|
{ std::lock_guard<std::mutex> lk(cv_mx_); }
|
|
cv_.notify_one();
|
|
}
|
|
|
|
std::size_t thread_count() const { return thread_count_; }
|
|
|
|
// ── IPoolProbe ────────────────────────────────────────────────────────────
|
|
|
|
PoolSnapshot snapshot(const std::string& name) const override {
|
|
std::size_t a = active_.load(std::memory_order_relaxed);
|
|
std::size_t t = total_.load(std::memory_order_relaxed);
|
|
return {
|
|
name, thread_count_,
|
|
t > a ? t - a : 0, // queued (approximate)
|
|
a, // executing
|
|
submitted_.load(std::memory_order_relaxed),
|
|
completed_.load(std::memory_order_relaxed),
|
|
};
|
|
}
|
|
|
|
private:
|
|
struct Task {
|
|
std::function<void()> fn;
|
|
float priority;
|
|
uint64_t seq;
|
|
// max-heap: higher priority runs first; older task wins tie
|
|
bool operator<(const Task& o) const {
|
|
if (priority != o.priority) return priority < o.priority;
|
|
return seq > o.seq;
|
|
}
|
|
};
|
|
|
|
// Separate cache lines to prevent false sharing between adjacent queues.
|
|
struct alignas(64) WorkerQueue {
|
|
std::priority_queue<Task> pq;
|
|
std::mutex mx;
|
|
};
|
|
|
|
std::optional<std::function<void()>> try_pop(WorkerQueue& q) {
|
|
std::lock_guard lock(q.mx);
|
|
if (q.pq.empty()) return std::nullopt;
|
|
auto fn = std::move(const_cast<Task&>(q.pq.top()).fn);
|
|
q.pq.pop();
|
|
return fn;
|
|
}
|
|
|
|
std::optional<std::function<void()>> try_steal(std::size_t thief) {
|
|
// Find the most-loaded peer without blocking — racy peek is fine.
|
|
std::size_t victim = thief, best = 0;
|
|
for (std::size_t i = 0; i < queues_.size(); ++i) {
|
|
if (i == thief) continue;
|
|
std::unique_lock lk(queues_[i]->mx, std::try_to_lock);
|
|
if (!lk) continue;
|
|
std::size_t n = queues_[i]->pq.size();
|
|
if (n > best) { best = n; victim = i; }
|
|
}
|
|
if (victim == thief) return std::nullopt;
|
|
return try_pop(*queues_[victim]);
|
|
}
|
|
|
|
void execute(std::function<void()>& fn) {
|
|
active_.fetch_add(1, std::memory_order_relaxed);
|
|
fn();
|
|
completed_.fetch_add(1, std::memory_order_relaxed);
|
|
active_.fetch_sub(1, std::memory_order_relaxed);
|
|
// Notify drain() if this was the last in-flight task.
|
|
// acq_rel ensures the decrement is visible before any drain() load.
|
|
// Lock drain_mx_ before notifying to avoid a lost wakeup against
|
|
// drain()'s predicate check (same hazard as submit()/cv_mx_).
|
|
if (total_.fetch_sub(1, std::memory_order_acq_rel) == 1) {
|
|
{ std::lock_guard<std::mutex> lk(drain_mx_); }
|
|
drain_cv_.notify_all();
|
|
}
|
|
}
|
|
|
|
void worker_loop(std::size_t id) {
|
|
while (true) {
|
|
if (auto fn = try_pop(*queues_[id])) { execute(*fn); continue; }
|
|
if (auto fn = try_steal(id)) { execute(*fn); continue; }
|
|
|
|
std::unique_lock lock(cv_mx_);
|
|
cv_.wait(lock, [this] {
|
|
return stopped_.load(std::memory_order_seq_cst)
|
|
|| total_.load(std::memory_order_relaxed) > 0;
|
|
});
|
|
if (stopped_.load(std::memory_order_seq_cst)
|
|
&& total_.load(std::memory_order_relaxed) == 0)
|
|
return;
|
|
}
|
|
}
|
|
|
|
const std::size_t thread_count_;
|
|
std::vector<std::unique_ptr<WorkerQueue>> queues_;
|
|
std::vector<std::thread> workers_;
|
|
|
|
std::mutex cv_mx_;
|
|
std::condition_variable cv_;
|
|
std::mutex drain_mx_;
|
|
std::condition_variable drain_cv_;
|
|
|
|
std::atomic<bool> stopped_{true};
|
|
std::atomic<size_t> total_{0}; // queued + executing
|
|
std::atomic<size_t> active_{0}; // executing only (for snapshot)
|
|
std::atomic<size_t> next_{0}; // round-robin submit cursor
|
|
std::atomic<uint64_t> seq_{0}; // tie-break for equal-priority tasks
|
|
std::atomic<uint64_t> submitted_{0};
|
|
std::atomic<uint64_t> completed_{0};
|
|
};
|
|
|
|
} // namespace kpn
|