KPN/include/kpn/shared_resource.hpp

#pragma once
#include "diagnostics.hpp"

#include <algorithm>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <functional>
#include <mutex>
#include <utility>
#include <vector>

namespace kpn {

template<typename T> class Channel; // forward declaration for acquire_balanced

// ── SharedResource ────────────────────────────────────────────────────────────
//
// Wraps an exclusive resource (e.g. an ONNX session, a CUDA stream) and
// arbitrates concurrent access using a priority-based waiter queue.
//
// When multiple nodes compete, the one with the highest priority score wins
// the next slot. Priority is re-evaluated at release time so it reflects the
// current queue state, not the state when the node first started waiting.
//
// Starvation prevention: each waiter's effective score grows with elapsed wait
// time (aging_per_second), ensuring a low-priority node eventually gets served.
//
// Usage:
//   SharedResource<OrtSession> res(session_args...);
//
//   // inside a node functor —
//   auto guard = res.acquire_balanced(in_channel, out_channel);
//   guard->Run(...);   // guard releases automatically on scope exit

template<typename T>
class SharedResource : public IResourceProbe {
public:
    // ── RAII guard ────────────────────────────────────────────────────────────

    class Guard {
        SharedResource* owner_;
        explicit Guard(SharedResource* o) : owner_(o) {}
        friend class SharedResource;
    public:
        Guard(Guard&& o) noexcept : owner_(std::exchange(o.owner_, nullptr)) {}
        Guard& operator=(Guard&&) = delete;
        Guard(const Guard&)       = delete;
        Guard& operator=(const Guard&) = delete;
        ~Guard() { if (owner_) owner_->release(); }

        T&  get()          { return owner_->resource_; }
        T*  operator->()   { return &owner_->resource_; }
        T&  operator*()    { return owner_->resource_; }
    };

    // ── Construction ──────────────────────────────────────────────────────────

    template<typename... Args>
    explicit SharedResource(Args&&... args)
        : resource_(std::forward<Args>(args)...) {}

    SharedResource(const SharedResource&)            = delete;
    SharedResource& operator=(const SharedResource&) = delete;
    SharedResource(SharedResource&&)                 = delete;
    SharedResource& operator=(SharedResource&&)      = delete;

    // ── Acquire ───────────────────────────────────────────────────────────────

    // Acquire with a callable that returns a priority in [0, 1].
    // Higher = more urgent. Called at every release to pick the best waiter.
    template<typename PriorityFn>
    Guard acquire(PriorityFn&& fn) {
        std::unique_lock lock(mutex_);
        if (!held_) {
            held_ = true;
            acq_.fetch_add(1, std::memory_order_relaxed);
            return Guard(this);
        }
        Waiter w{std::function<float()>(std::forward<PriorityFn>(fn)), clock_t::now()};
        waiters_.push_back(&w);
        update_peak(waiters_.size());
        current_waiters_.store(waiters_.size(), std::memory_order_relaxed);

        auto t0 = w.wait_start;
        w.cv.wait(lock, [&w] { return w.ready; });

        int64_t wait_us = std::chrono::duration_cast<std::chrono::microseconds>(
            clock_t::now() - t0).count();
        waiters_.erase(std::find(waiters_.begin(), waiters_.end(), &w));
        current_waiters_.store(waiters_.size(), std::memory_order_relaxed);
        acq_.fetch_add(1, std::memory_order_relaxed);
        total_wait_us_.fetch_add(static_cast<uint64_t>(wait_us > 0 ? wait_us : 0),
                                 std::memory_order_relaxed);
        return Guard(this);
    }

    // Acquire with no priority (all waiters treated equally, order is fair-ish).
    Guard acquire() {
        return acquire([] { return 0.5f; });
    }

    // Acquire with priority derived from channel fill fractions:
    //   score = input_fill × output_headroom
    // A node with a full input queue and empty output queue has the highest
    // urgency — it has work to do and nowhere to stall downstream.
    template<typename In, typename Out>
    Guard acquire_balanced(const Channel<In>& in_ch, const Channel<Out>& out_ch) {
        return acquire([&in_ch, &out_ch] {
            float in_fill  = in_ch.capacity()  ? float(in_ch.size())  / in_ch.capacity()  : 0.5f;
            float out_head = out_ch.capacity() ? 1.0f - float(out_ch.size()) / out_ch.capacity() : 0.5f;
            return in_fill * out_head;
        });
    }

    // ── IResourceProbe ────────────────────────────────────────────────────────

    ResourceSnapshot snapshot(const std::string& name) const override {
        std::lock_guard lock(mutex_);
        uint64_t a   = acq_.load(std::memory_order_relaxed);
        uint64_t w   = total_wait_us_.load(std::memory_order_relaxed);
        return {
            name,
            a,
            a > 0 ? double(w) / a / 1000.0 : 0.0,
            peak_waiters_.load(std::memory_order_relaxed),
            current_waiters_.load(std::memory_order_relaxed),
            held_,
        };
    }

private:
    void release() {
        std::unique_lock lock(mutex_);
        if (waiters_.empty()) {
            held_ = false;
            return;
        }
        // Re-evaluate every waiter's current priority and apply aging bonus.
        auto now = clock_t::now();
        Waiter* best       = nullptr;
        float   best_score = -1.0f;
        for (Waiter* w : waiters_) {
            float age_s = std::chrono::duration<float>(now - w->wait_start).count();
            float score = w->priority_fn() + age_s * kAgingPerSecond;
            if (score > best_score) { best_score = score; best = w; }
        }
        best->ready = true;
        best->cv.notify_one();
        // held_ stays true — ownership transfers to the woken waiter.
    }

    void update_peak(std::size_t n) {
        uint64_t prev = peak_waiters_.load(std::memory_order_relaxed);
        while (n > prev &&
               !peak_waiters_.compare_exchange_weak(prev, n,
                   std::memory_order_relaxed, std::memory_order_relaxed))
            ;
    }

    struct Waiter {
        std::function<float()>  priority_fn;
        clock_t::time_point     wait_start;
        std::condition_variable cv;
        bool                    ready{false};

        Waiter(std::function<float()> fn, clock_t::time_point t)
            : priority_fn(std::move(fn)), wait_start(t) {}
    };

    static constexpr float kAgingPerSecond = 0.05f;

    T                        resource_;
    bool                     held_{false};
    mutable std::mutex       mutex_;
    std::vector<Waiter*>     waiters_;
    std::atomic<uint64_t>    acq_{0};
    std::atomic<uint64_t>    total_wait_us_{0};
    std::atomic<uint64_t>    peak_waiters_{0};
    std::atomic<uint64_t>    current_waiters_{0};
};

// ── Factory ───────────────────────────────────────────────────────────────────

template<typename T, typename... Args>
SharedResource<T> make_shared_resource(Args&&... args) {
    return SharedResource<T>(std::forward<Args>(args)...);
}

} // namespace kpn