dtourolle
771b9f8593
abbb2d4guarded submit() against stop() with a shared/exclusive lock, because stop() ends with queues_.clear() and submit() indexes queues_. It missed the other writer: start() rebuilds the same vector and took no lock at all. A submission can genuinely land while a pool is inside start(). A network starts its nodes one at a time, and a node already started fires into the next one's channel, whose push callback submits. ThreadSanitizer reports it as the read at scheduler.hpp:114 against the write at :59, and the consequence is worse than a torn read: push_back can reallocate the vector under a reader that has already indexed it. Surfaced by6802328. Firing push_callback_ unconditionally is correct — the argument in that commit holds — and it makes callbacks frequent enough during startup to hit this window. It went from unobserved to 4 races in one run of the unit suite. Holding the lock across the thread spawn is safe: all queues are constructed before any worker starts, and worker_loop never takes the lifecycle lock, so there is nothing for it to deadlock against. Verified with -DKPN_SANITIZER=thread: 4 races in one run of five before, 0 across five unit runs and two stress runs after. 148/148.
280 lines
12 KiB
C++
280 lines
12 KiB
C++
#pragma once
|
|
#include "diagnostics.hpp"
|
|
#include <atomic>
|
|
#include <condition_variable>
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <shared_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 {
|
|
// Under the lifecycle lock for the same reason stop() is: submit()
|
|
// reads queues_ and this rebuilds it. A network starts its nodes one at
|
|
// a time, and a node already started fires into the next one's channel,
|
|
// whose push callback submits — so a submission can genuinely land
|
|
// while another pool is still inside start(). ThreadSanitizer reports
|
|
// it as a read at submit() against this write, and the consequence is
|
|
// worse than a torn read: push_back can reallocate the vector under a
|
|
// reader that has already indexed it.
|
|
//
|
|
// Queues are all constructed before any worker is spawned, which is
|
|
// what keeps worker_loop's own queues_[id] out of this — it never takes
|
|
// the lock, so holding it across the spawn cannot deadlock.
|
|
std::unique_lock lk(lifecycle_mx_);
|
|
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 {
|
|
// Close the pool to new work before touching anything, and do it under
|
|
// the lifecycle lock so no submit() is midway through indexing queues_.
|
|
{
|
|
std::unique_lock lk(lifecycle_mx_);
|
|
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);
|
|
queued_.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();
|
|
// Join without the lock: a worker's task may call submit(), which takes
|
|
// it shared, and holding it here would deadlock against that.
|
|
for (auto& t : workers_) if (t.joinable()) t.join();
|
|
// Destroying the queues is what submit() must never race. By now
|
|
// stopped_ is published, so any submit() that acquires the lock after
|
|
// this point returns without touching them.
|
|
std::unique_lock lk(lifecycle_mx_);
|
|
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 {
|
|
// A submission can arrive after this pool has been stopped, and did so
|
|
// by an ordinary route: a node's space callback fires from whichever
|
|
// thread drained the channel, which belongs to the *consumer*. Stop the
|
|
// producer first — as a sources-first shutdown does — and the consumer
|
|
// keeps draining its backlog, firing the producer's space callback into
|
|
// a pool whose stop() has already run queues_.clear(). submit() then
|
|
// indexed an empty vector: a segfault, reproducible about 12 runs in 20.
|
|
//
|
|
// The shared lock is what makes the check meaningful. Reading stopped_
|
|
// alone leaves the window between the read and the indexing, which is
|
|
// precisely where stop() clears the vector.
|
|
std::shared_lock lk(lifecycle_mx_);
|
|
if (stopped_.load(std::memory_order_acquire) || queues_.empty()) {
|
|
rejected_.fetch_add(1, std::memory_order_relaxed);
|
|
return;
|
|
}
|
|
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);
|
|
queued_.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_; }
|
|
|
|
/// Submissions dropped because the pool was stopped. See rejected_.
|
|
uint64_t rejected() const { return rejected_.load(std::memory_order_relaxed); }
|
|
|
|
// ── IPoolProbe ────────────────────────────────────────────────────────────
|
|
|
|
PoolSnapshot snapshot(const std::string& name) const override {
|
|
std::size_t a = active_.load(std::memory_order_relaxed);
|
|
return {
|
|
name, thread_count_,
|
|
queued_.load(std::memory_order_relaxed), // queued (exact)
|
|
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();
|
|
queued_.fetch_sub(1, std::memory_order_relaxed);
|
|
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)
|
|
|| queued_.load(std::memory_order_relaxed) > 0;
|
|
});
|
|
// Exit on queued_, not total_: waiting for total_ to reach zero
|
|
// meant waiting for someone else's task to finish, which this
|
|
// worker cannot help with and would spin through until it did.
|
|
if (stopped_.load(std::memory_order_seq_cst)
|
|
&& queued_.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_;
|
|
|
|
/// Guards the lifetime of queues_/workers_ against a concurrent submit().
|
|
/// Shared by submit, exclusive by stop, so submissions still run in
|
|
/// parallel with each other.
|
|
mutable std::shared_mutex lifecycle_mx_;
|
|
|
|
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 (drain() waits on this)
|
|
/// Queued only — never counts a task that is already executing.
|
|
///
|
|
/// The wait predicate used total_, which includes running tasks, so while
|
|
/// any one task ran every *other* worker's predicate was true: wait()
|
|
/// returned instantly and the worker spun through try_pop / try_steal /
|
|
/// wait at full speed, try_lock-ing every peer queue on each pass. One slow
|
|
/// task therefore pinned every other core and contended the very mutexes
|
|
/// the working thread needed. Sleeping requires "no work is *waiting*",
|
|
/// which is this.
|
|
std::atomic<size_t> queued_{0}; // waiting to run
|
|
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};
|
|
/// Submissions refused because the pool was already stopped. Not an error —
|
|
/// teardown races are expected — but silence here would hide a node that
|
|
/// keeps trying to run after its pool is gone.
|
|
std::atomic<uint64_t> rejected_{0};
|
|
std::atomic<uint64_t> completed_{0};
|
|
};
|
|
|
|
} // namespace kpn
|