This commit is contained in:
@@ -0,0 +1,190 @@
|
||||
#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
|
||||
Reference in New Issue
Block a user