dtourolle
7b7f631e6d
SharedResource::acquire() blocks on a condition variable whose predicate only becomes true when release() hands over ownership. No timeout, no stop condition. A node parked there was not observing stop flags, so teardown had no way to reach it: the worker never returned, the pool's join never completed, and shutdown waited on a resource nobody was going to release — which is precisely the situation when the holder is being stopped too. close() wakes every waiter and refuses further acquisitions, and the waiters leave through ResourceClosedError, which is an exception the node error path already handles rather than a new mechanism. StaticNetwork calls it on registered resources at the top of halt() and shutdown(), before stopping any node, since a node stopped while parked cannot respond to being stopped. The handover needed care in two places. A waiter woken by close() has not been given ownership, so it takes no Guard and leaves held_ exactly as it found it; and release() now skips handing over to waiters when closed, because handing ownership to a thread that is on its way out would leave held_ true with nobody holding it. reopen() is there for reuse across runs, which the persistent-pipeline work will want; teardown does not need it. Verified in both directions: without close() the waiter thread never returns and the test's join blocks; with it the waiter leaves through ResourceClosedError while the holder still has the resource. 145/145.
237 lines
9.8 KiB
C++
237 lines
9.8 KiB
C++
#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
|
||
|
||
/// Thrown by a pending acquire() when the resource is closed underneath it.
|
||
///
|
||
/// acquire() blocks on a condition variable with no timeout and no stop
|
||
/// condition, so a node parked there ignored teardown entirely: the worker
|
||
/// never returned, the pool's join never completed, and shutdown hung on a
|
||
/// resource nobody was going to release. Closing the resource turns that into
|
||
/// an exception the node's normal error path already handles.
|
||
class ResourceClosedError : public std::runtime_error {
|
||
public:
|
||
ResourceClosedError() : std::runtime_error("shared resource closed") {}
|
||
};
|
||
|
||
// ── 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 (closed_) throw ResourceClosedError{};
|
||
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;
|
||
// Woken either by release() handing over ownership, or by close()
|
||
// giving up on the wait entirely.
|
||
w.cv.wait(lock, [&w] { return w.ready || w.closed; });
|
||
|
||
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);
|
||
total_wait_us_.fetch_add(static_cast<uint64_t>(wait_us > 0 ? wait_us : 0),
|
||
std::memory_order_relaxed);
|
||
|
||
// Closed without being handed ownership: no Guard, so nothing to
|
||
// release, and held_ is left exactly as close() found it.
|
||
if (!w.ready) throw ResourceClosedError{};
|
||
|
||
acq_.fetch_add(1, std::memory_order_relaxed);
|
||
return Guard(this);
|
||
}
|
||
|
||
/// Wake every waiter and refuse further acquisitions.
|
||
///
|
||
/// Teardown is the whole point: a node parked in acquire() is not
|
||
/// observing stop flags, so without this the only way out is for whoever
|
||
/// holds the resource to release it — which, if that node is also being
|
||
/// stopped, may never happen. Idempotent, and safe to call from any thread.
|
||
void close() override {
|
||
std::lock_guard lock(mutex_);
|
||
closed_ = true;
|
||
for (Waiter* w : waiters_) {
|
||
w->closed = true;
|
||
w->cv.notify_one();
|
||
}
|
||
}
|
||
|
||
/// Reopen after a close(). For reuse across runs; not needed for teardown.
|
||
void reopen() {
|
||
std::lock_guard lock(mutex_);
|
||
closed_ = false;
|
||
}
|
||
|
||
// 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_);
|
||
// Hand over only to a waiter that is still waiting. A closed one is on
|
||
// its way out and will not take ownership, so treating it as the next
|
||
// holder would leave held_ true with nobody holding it.
|
||
if (closed_ || 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}; // handed ownership by release()
|
||
bool closed{false}; // woken by close() instead
|
||
|
||
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};
|
||
bool closed_{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
|