KPN/include/kpn/scheduler.hpp

#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 {
        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