diff --git a/src/runtime/Cargo.toml b/src/runtime/Cargo.toml index 007046f2..d8457ab1 100644 --- a/src/runtime/Cargo.toml +++ b/src/runtime/Cargo.toml @@ -19,3 +19,7 @@ path = "../../tests/runtime/scheduler.rs" [[test]] name = "engine" path = "../../tests/runtime/engine.rs" + +[[test]] +name = "admission" +path = "../../tests/runtime/admission.rs" diff --git a/src/runtime/src/admission.rs b/src/runtime/src/admission.rs new file mode 100644 index 00000000..8a96f5ae --- /dev/null +++ b/src/runtime/src/admission.rs @@ -0,0 +1,219 @@ +//! Optional capacity limits for a worker that wants them. +//! +//! `SerialScheduler` admits one unit of work at a time and never refuses: a hundred requests +//! queue behind each other and the only bound is the caller's patience. That is the right +//! default for the workers here today, and this type is what a worker opts into when it wants a +//! bound instead. +//! +//! **Units are the caller's, and they are explicit.** Admission counts two things: how many +//! requests are outstanding, and how much each one costs. A request counts as one; the cost is +//! whatever the caller passes, so a prefill worker can admit by token count, a worker with +//! large payloads by bytes, and a worker that cares about neither by `1`. Nothing here guesses +//! which of those a model means. +//! +//! **The permit is owned, not borrowed.** It holds a share of the admission, so it can be moved +//! into the blocking work it was granted for — which is what the runtime's executor already +//! does — and the capacity is released when that work finishes with it rather than when the +//! request that asked for it goes away. +//! +//! This is admission only. It decides whether work may start and hands back a permit that +//! releases when the work is done with; it does not own a thread, run anything, or move a +//! request between states. Readiness, worker-thread ownership and drain are separate. + +use std::sync::{Arc, Mutex}; + +use crate::engine::Report; + +/// A bound a worker may opt into. Clones share it, so a queue can hold one and so can the +/// handlers that feed it. +/// +/// ``` +/// use omni_runtime::admission::BoundedAdmission; +/// +/// let admission = BoundedAdmission::new(2, None); +/// let first = admission.try_admit(1).expect("one of two"); +/// let second = admission.try_admit(1).expect("two of two"); +/// assert!(admission.try_admit(1).is_none(), "the third has nowhere to go"); +/// +/// // A permit releases its capacity when the work is done with. +/// drop(first); +/// assert!(admission.try_admit(1).is_some()); +/// ``` +/// +/// It wraps a shared inner rather than *being* an [`Arc`] alias: an alias would inherit `Arc`'s +/// associated functions, so `BoundedAdmission::new` would resolve to `Arc::new`. +#[derive(Clone, Debug)] +pub struct BoundedAdmission(Arc); + +/// The bound and the counters behind one [`BoundedAdmission`]. +#[derive(Debug)] +struct Admission { + /// Requests that may be outstanding at once. Zero would admit nothing, so it is refused at + /// construction rather than silently rejecting every request. + requests: usize, + /// Total cost that may be outstanding, when the caller wants to bound something other than + /// a request count. `None` bounds requests only. + units: Option, + held: Mutex, +} + +/// What admission is currently holding, and what it has refused. +#[derive(Debug, Default)] +struct Held { + outstanding: usize, + units: usize, + rejected: u64, +} + +/// One admitted piece of work's reservation, released when its [`Permit`] drops. +/// +/// `units` preserves the caller's original cost, as reported by [`Permit::units`]. `charged` +/// records the amount added to the shared unit total: the same cost when a unit budget is +/// configured, or zero when only requests are bounded. Release subtracts exactly `charged`. +/// Requests exceeding the unit ceiling are refused rather than clamped. +#[derive(Clone, Copy, Debug)] +pub struct Reservation { + /// The cost the work was admitted at, which is what [`Permit::units`] reports. + units: usize, + /// What the shared unit total was advanced by, so release subtracts exactly that. Zero when + /// no unit budget is configured: there is then no total, and charging it would overflow for + /// a cost near `usize::MAX` while bounding nothing. + charged: usize, +} + +impl BoundedAdmission { + /// A bound of `requests` at a time, optionally also of `units` at a time. + /// + /// # Panics + /// + /// If `requests` is zero, or if `units` is `Some(0)`. Both would admit nothing at all, + /// which is a configuration mistake rather than a limit worth enforcing, and it would show + /// up as every request refused rather than as the mistake it is. + pub fn new(requests: usize, units: Option) -> Self { + assert!( + requests > 0, + "admission with no request capacity admits nothing" + ); + assert!( + units != Some(0), + "admission with no unit capacity admits nothing" + ); + Self(Arc::new(Admission { + requests, + units, + held: Mutex::new(Held::default()), + })) + } + + /// Admits one piece of work costing `units`, or refuses it because there is no room. + /// + /// The permit is the reservation: capacity is held until it is dropped, which is what + /// makes "outstanding" mean "accepted and not yet finished with" rather than "accepted". + /// A refusal is counted and nothing is reserved. + /// + /// A request costing more than the unit ceiling is refused rather than admitted and + /// overrunning it, so a reservation is never larger than the ceiling and the held total + /// cannot reach `usize::MAX` by accumulation. + /// + /// The permit owns a share of this admission, so it can outlive the handle that issued it + /// and be moved into the blocking work it was granted for. + pub fn try_admit(&self, units: usize) -> Option { + let mut held = self.0.lock(); + if held.outstanding >= self.0.requests { + held.rejected += 1; + return None; + } + let charged = match self.0.units { + Some(ceiling) => { + // A request larger than the whole ceiling is refused rather than clamped, so a + // reservation is never bigger than the budget it was admitted under. What is + // left of the ceiling is then known to fit, because `held.units <= ceiling` + // holds as an invariant once every reservation is bounded this way. + if units > ceiling || held.units > ceiling - units { + held.rejected += 1; + return None; + } + units + } + // No unit budget, so there is no total to keep and nothing is charged. Adding the + // caller's cost here would overflow for a cost near `usize::MAX` while protecting + // nothing: the request bound is the only limit this configuration asked for. + None => 0, + }; + held.outstanding += 1; + held.units += charged; + Some(Permit { + admission: Arc::clone(&self.0), + reserved: Reservation { units, charged }, + }) + } + + /// What a worker would publish as its queue state. `capacity` is the request bound, which + /// is the one every caller of this type has. + pub fn report(&self) -> Report { + let held = self.0.lock(); + Report { + depth: held.outstanding, + capacity: self.0.requests, + rejected: held.rejected, + } + } +} + +impl Admission { + /// Releases one reservation. Called when its [`Permit`] drops. + /// + /// Guarded rather than saturating: a reservation larger than what is held would mean two + /// permits for one admission, and quietly underflowing would hide that. The public API + /// cannot produce one, because a reservation is refused unless it fits under the ceiling. + fn release(&self, reservation: Reservation) { + let mut held = self.lock(); + held.outstanding = held + .outstanding + .checked_sub(1) + .expect("a permit releases exactly one admission"); + if self.units.is_some() { + held.units = held + .units + .checked_sub(reservation.charged) + .expect("a permit releases the units it was admitted with"); + } + } + + fn lock(&self) -> std::sync::MutexGuard<'_, Held> { + // A panic while holding this leaves the counters readable and consistent, so a + // poisoned lock is not a reason to refuse every later request. + self.held.lock().unwrap_or_else(|error| error.into_inner()) + } +} + +/// One admitted piece of work. Holding it holds its capacity; dropping it releases both. +/// +/// It owns a share of its admission rather than borrowing it, so it can be moved into the +/// blocking work it was granted for. The runtime's own executor does exactly that; a borrow +/// would not compile there, because that work has to be `'static`. +/// +/// This is deliberately not `Clone` and not `Copy`: two permits for one admission would release +/// capacity twice, which is the whole failure this type exists to prevent. +#[derive(Debug)] +pub struct Permit { + admission: Arc, + reserved: Reservation, +} + +impl Permit { + /// The cost this work was admitted at, for a caller that wants to report or account it. + /// + /// This is the cost the caller passed, unchanged: nothing is clamped, because a request too + /// large for the unit ceiling is refused rather than shaved to fit. When no unit budget is + /// configured the cost is still reported, even though no unit total was kept for it. + pub fn units(&self) -> usize { + self.reserved.units + } +} + +impl Drop for Permit { + fn drop(&mut self) { + self.admission.release(self.reserved); + } +} diff --git a/src/runtime/src/lib.rs b/src/runtime/src/lib.rs index dde9174f..c3dc515c 100644 --- a/src/runtime/src/lib.rs +++ b/src/runtime/src/lib.rs @@ -1,5 +1,6 @@ //! Shared admission, dispatch and the worker-side engine contract for a loaded executor. +pub mod admission; pub mod engine; use std::sync::Arc; diff --git a/tests/runtime/admission.rs b/tests/runtime/admission.rs new file mode 100644 index 00000000..ea47b86c --- /dev/null +++ b/tests/runtime/admission.rs @@ -0,0 +1,331 @@ +//! Optional capacity limits: what they admit, what they refuse, and what a permit means. +//! +//! Everything goes through the public API, because that is all a worker sees. + +use std::{ + sync::{ + Arc, Barrier, + atomic::{AtomicUsize, Ordering}, + }, + thread, +}; + +use omni_runtime::admission::BoundedAdmission; + +#[test] +fn a_request_bound_refuses_once_it_is_full() { + let admission = BoundedAdmission::new(2, None); + let first = admission.try_admit(1).expect("the first of two"); + let second = admission.try_admit(1).expect("the second of two"); + assert!( + admission.try_admit(1).is_none(), + "the third has nowhere to go" + ); + + let report = admission.report(); + assert_eq!(report.depth, 2, "a refused request must not be counted"); + assert_eq!(report.capacity, 2); + assert_eq!(report.rejected, 1); + + // A refusal is not a lost slot: releasing one makes room for the next. + drop(first); + assert_eq!(admission.report().depth, 1); + assert!(admission.try_admit(1).is_some()); + drop(second); +} + +/// The point of a permit: capacity means "accepted and not yet finished with". A worker that +/// hands work to a thread and forgets to release would be bounded forever, so release is tied +/// to the permit's lifetime rather than to a call someone has to remember. +#[test] +fn capacity_is_held_until_the_permit_is_dropped() { + let admission = BoundedAdmission::new(1, None); + let permit = admission.try_admit(1).expect("one of one"); + assert!(admission.try_admit(1).is_none(), "still held"); + drop(permit); + assert_eq!(admission.report().depth, 0); + assert!(admission.try_admit(1).is_some(), "released"); +} + +/// Units are the caller's, which is what lets a prefill worker bound tokens and a worker with +/// large payloads bound bytes without this type knowing which it is looking at. +#[test] +fn a_unit_bound_admits_by_cost_rather_than_by_count() { + // Room for three requests, or 100 units, whichever runs out first. + let admission = BoundedAdmission::new(3, Some(100)); + + let big = admission.try_admit(60).expect("60 of 100 units"); + assert_eq!(admission.report().depth, 1, "one request, 60 units"); + + let small = admission.try_admit(40).expect("the remaining 40 units"); + assert!( + admission.try_admit(1).is_none(), + "the unit ceiling is reached even though only two requests are outstanding" + ); + + // Releasing the large one frees its units, not just its request slot. + drop(big); + assert!(admission.try_admit(50).is_some(), "60 units came back"); + drop(small); +} + +/// A request that costs more than the whole ceiling is refused rather than admitted and +/// overrunning it — the case a naive `depth < capacity` check gets wrong. +#[test] +fn one_request_larger_than_the_ceiling_is_refused() { + let admission = BoundedAdmission::new(10, Some(100)); + assert!(admission.try_admit(101).is_none()); + assert_eq!(admission.report().depth, 0, "nothing was reserved"); + assert!( + admission.try_admit(100).is_some(), + "exactly the ceiling fits" + ); +} + +#[test] +fn the_request_count_is_reported_as_capacity() { + let admission = BoundedAdmission::new(7, Some(1_000)); + assert_eq!(admission.report().capacity, 7); + assert_eq!(admission.report().depth, 0); + assert_eq!(admission.report().rejected, 0); +} + +/// A zero bound admits nothing, which is a configuration mistake rather than a limit. It is +/// refused where it is written instead of showing up as every request being turned away. +#[test] +#[should_panic(expected = "admits nothing")] +fn a_zero_request_bound_is_refused_at_construction() { + let _ = BoundedAdmission::new(0, None); +} + +#[test] +#[should_panic(expected = "admits nothing")] +fn a_zero_unit_bound_is_refused_at_construction() { + let _ = BoundedAdmission::new(1, Some(0)); +} + +/// Clones share the bound, so the queue and the handlers feeding it cannot disagree about how +/// much room is left. +#[test] +fn clones_share_one_bound() { + let admission = BoundedAdmission::new(1, None); + let shared = admission.clone(); + let permit = shared.try_admit(1).expect("one of one"); + assert!(admission.try_admit(1).is_none(), "the clone held the slot"); + drop(permit); + assert!(admission.try_admit(1).is_some()); + assert_eq!(admission.report().rejected, 1, "both see the same count"); +} + +/// Regression: every admission is released by exactly what it reserved, in any order. The +/// permit stores the units it actually reserved rather than the units it was asked for, so the +/// counter cannot be driven negative — an earlier version stored the requested cost and, where +/// the two differed, released past zero and panicked on a later release. +#[test] +fn a_reservation_is_released_by_what_was_reserved_not_by_what_was_asked() { + // Room for many requests and exactly 100 units, so the unit ceiling is what binds. + let admission = BoundedAdmission::new(10, Some(100)); + + let a = admission.try_admit(10).expect("10 of 100"); + let b = admission.try_admit(60).expect("60 more"); + let c = admission.try_admit(30).expect("the remaining 30"); + assert!( + admission.try_admit(1).is_none(), + "the ceiling is exactly reached" + ); + + // Release order is irrelevant when what goes in is what comes out. + drop(b); + let again = admission.try_admit(60).expect("b's units came back"); + drop(a); + drop(c); + assert_eq!( + admission.report().depth, + 1, + "only the re-admitted 60 is held" + ); + drop(again); + assert_eq!(admission.report().depth, 0, "and it released its own units"); + assert!( + admission.try_admit(100).is_some(), + "the ceiling is whole again" + ); +} + +/// A request costing more than the ceiling is refused outright, which is what keeps the held +/// total below `usize::MAX`: a reservation larger than the ceiling it was admitted under would +/// be released against the others and unbalance the counter. +#[test] +fn a_request_larger_than_the_unit_ceiling_is_refused_however_large_it_is() { + let admission = BoundedAdmission::new(10, Some(64)); + + for cost in [65, 1_000, usize::MAX] { + assert!( + admission.try_admit(cost).is_none(), + "a request costing {cost} must be refused, not clamped" + ); + } + assert_eq!(admission.report().depth, 0, "nothing was reserved"); + assert_eq!(admission.report().rejected, 3, "all three were counted"); + + // The largest request that fits is still admitted, and released with its own units. + let full = admission.try_admit(64).expect("exactly the ceiling"); + assert_eq!(full.units(), 64); + assert!( + admission.try_admit(1).is_none(), + "and it took the whole ceiling" + ); + drop(full); + assert_eq!(admission.report().depth, 0); + assert!(admission.try_admit(64).is_some(), "the ceiling is reusable"); +} + +/// The ceiling is the whole budget, so one request may take all of it and no more. +#[test] +fn the_ceiling_is_the_largest_single_reservation() { + let admission = BoundedAdmission::new(10, Some(usize::MAX)); + + let all = admission.try_admit(usize::MAX).expect("the whole ceiling"); + assert!( + admission.try_admit(1).is_none(), + "even one unit is refused: the ceiling is taken" + ); + drop(all); + assert_eq!(admission.report().depth, 0); + assert!( + admission.try_admit(usize::MAX).is_some(), + "and the whole ceiling comes back" + ); +} + +/// Concurrent admitters can never together hold more than the request bound, and every +/// admission is released exactly once. A counter that drifted one way would let the live count +/// pass the bound; drifting the other way would leave permits outstanding at the end. +#[test] +fn concurrent_admissions_never_exceed_the_bound_and_all_come_back() { + const THREADS: usize = 32; + const ROUNDS: usize = 100; + const CAPACITY: usize = 4; + + for _ in 0..ROUNDS { + let admission = BoundedAdmission::new(CAPACITY, None); + let barrier = Arc::new(Barrier::new(THREADS)); + let live = Arc::new(AtomicUsize::new(0)); + let peak = Arc::new(AtomicUsize::new(0)); + + let handles: Vec<_> = (0..THREADS) + .map(|_| { + let admission = admission.clone(); + let barrier = Arc::clone(&barrier); + let live = Arc::clone(&live); + let peak = Arc::clone(&peak); + thread::spawn(move || { + barrier.wait(); + if let Some(_permit) = admission.try_admit(1) { + let now = live.fetch_add(1, Ordering::SeqCst) + 1; + peak.fetch_max(now, Ordering::SeqCst); + live.fetch_sub(1, Ordering::SeqCst); + } + }) + }) + .collect(); + for handle in handles { + handle.join().expect("no admitting thread may panic"); + } + + assert_eq!( + live.load(Ordering::SeqCst), + 0, + "every admission was released" + ); + assert!( + peak.load(Ordering::SeqCst) <= CAPACITY, + "at most the capacity is ever held at once" + ); + assert_eq!(admission.report().depth, 0, "every admission was released"); + assert_eq!( + admission.report().capacity, + CAPACITY, + "the bound outlives the contention" + ); + } +} + +/// A panic while the counters are locked must not refuse every later request: the counters are +/// still consistent, so the lock is recovered rather than poisoned for good. +#[test] +fn a_poisoned_lock_keeps_admitting() { + let admission = BoundedAdmission::new(1, None); + + let held = admission.clone(); + let panicked = thread::spawn(move || { + let _permit = held.try_admit(1).expect("one of one"); + panic!("admit and then die holding the lock"); + }) + .join(); + assert!(panicked.is_err(), "the helper thread was meant to panic"); + + let report = admission.report(); + assert_eq!( + report.depth, 0, + "the panicking thread's permit was dropped as it unwound, so its slot came back" + ); + assert!( + admission.try_admit(1).is_some(), + "the recovered lock admits" + ); +} + +/// A permit is owned, so it can be moved into the work it was granted for instead of being tied +/// to the task that asked for it. This is the shape the runtime's executor uses, and the reason +/// the permit is not a borrow of the admission. +#[test] +fn a_permit_can_be_moved_into_the_work_it_was_granted_for() { + let admission = BoundedAdmission::new(1, None); + let permit = admission.try_admit(1).expect("one of one"); + assert!(admission.try_admit(1).is_none(), "held by the permit"); + + let held = admission.clone(); + thread::spawn(move || { + // The work owns the permit and releases the capacity when it finishes with it. + let _permit = permit; + assert!( + held.try_admit(1).is_none(), + "still held while the work runs" + ); + }) + .join() + .expect("the work thread may not panic"); + + assert_eq!(admission.report().depth, 0, "the work released its slot"); + assert!(admission.try_admit(1).is_some(), "and it can be reused"); +} + +/// Regression: a request-only bound has no unit budget, so nothing is accumulated. Adding the +/// caller's cost anyway overflowed for a cost near `usize::MAX`: debug builds panicked on +/// admission, release builds wrapped the total and then panicked in `release`. +#[test] +fn a_request_only_bound_does_not_accumulate_units() { + for reverse in [false, true] { + let admission = BoundedAdmission::new(2, None); + let huge = admission.try_admit(usize::MAX).expect("one of two"); + let one = admission.try_admit(1).expect("two of two"); + assert_eq!(admission.report().depth, 2); + assert_eq!(huge.units(), usize::MAX, "the caller's cost is preserved"); + assert_eq!(one.units(), 1, "and so is the small one's"); + + // Both drop orders have to release cleanly and leave the same depth. + if reverse { + drop(one); + drop(huge); + } else { + drop(huge); + drop(one); + } + assert_eq!(admission.report().depth, 0, "every admission came back"); + assert!( + admission.try_admit(usize::MAX).is_some(), + "and the bound admits again afterwards" + ); + } +}