Performance improvements, better readme and complete python bindings

include/kpn/scheduler.hpp

@@ -0,0 +1,198 @@
+#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);
+        }
+        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);
+        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.
+        if (total_.fetch_sub(1, std::memory_order_acq_rel) == 1)
+            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