use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::{Mutex, PoisonError};
use tokio::sync::{Notify, Semaphore, SemaphorePermit};
use tokio_util::sync::CancellationToken;
pub const SILENT_YIELDS_MAX: u32 = 3;
struct SilentOpen {
label: &'static str,
token: CancellationToken,
tokens: u64,
yields: bool,
}
pub struct SessionBudget {
permits: Semaphore,
silent_permits: Semaphore,
silent_open: Mutex<Option<SilentOpen>>,
pool: Option<u64>,
in_flight: Mutex<u64>,
room: Notify,
density: [AtomicU64; Shape::COUNT],
displacing: AtomicUsize,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Shape {
Turn,
Run,
Loop,
Roll,
Title,
Impersonation,
Summary,
}
impl Shape {
pub const COUNT: usize = 7;
}
pub struct Reservation<'a> {
budget: &'a SessionBudget,
tokens: u64,
label: Option<&'static str>,
cancel: CancellationToken,
child: Option<CancellationToken>,
_permit: SemaphorePermit<'a>,
}
impl Reservation<'_> {
pub fn stream_token(&self) -> CancellationToken {
self.child.clone().unwrap_or_else(|| self.cancel.clone())
}
pub fn displaced(&self) -> bool {
self.child.as_ref().is_some_and(|c| c.is_cancelled()) && !self.cancel.is_cancelled()
}
}
impl Drop for Reservation<'_> {
fn drop(&mut self) {
if let Some(label) = self.label {
let mut silent = self.budget.silent();
if silent.as_ref().is_some_and(|s| s.label == label) {
*silent = None;
}
}
if self.tokens > 0 {
let mut open = self.budget.open();
*open = open.saturating_sub(self.tokens);
drop(open);
self.budget.room.notify_waiters();
}
}
}
struct Displacing<'a>(&'a SessionBudget);
impl<'a> Displacing<'a> {
fn new(budget: &'a SessionBudget) -> Self {
budget.displacing.fetch_add(1, Ordering::SeqCst);
Self(budget)
}
}
impl Drop for Displacing<'_> {
fn drop(&mut self) {
self.0.displacing.fetch_sub(1, Ordering::SeqCst);
self.0.room.notify_waiters();
}
}
impl std::fmt::Debug for SessionBudget {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SessionBudget")
.field("sessions", &self.permits.available_permits())
.field("pool", &self.pool)
.field("in_flight", &*self.open())
.field(
"density",
&self
.density
.iter()
.map(|d| match d.load(Ordering::Relaxed) {
0 => 1.0,
bits => f64::from_bits(bits).max(1.0),
})
.collect::<Vec<_>>(),
)
.field("displacing", &self.displacing.load(Ordering::SeqCst))
.finish()
}
}
impl SessionBudget {
pub fn new(sessions: u32, pool: Option<u64>) -> Self {
Self {
permits: Semaphore::new(sessions.max(1) as usize),
silent_permits: Semaphore::new(1),
silent_open: Mutex::new(None),
pool: pool.filter(|&n| n > 0),
in_flight: Mutex::new(0),
room: Notify::new(),
density: std::array::from_fn(|_| AtomicU64::new(0)),
displacing: AtomicUsize::new(0),
}
}
#[cfg(test)]
pub fn pool(&self) -> Option<u64> {
self.pool
}
#[cfg(test)]
pub fn available_sessions(&self) -> usize {
self.permits.available_permits()
}
#[cfg(test)]
pub fn in_flight(&self) -> u64 {
*self.open()
}
#[cfg(test)]
pub fn displacing(&self) -> usize {
self.displacing.load(Ordering::SeqCst)
}
fn open(&self) -> std::sync::MutexGuard<'_, u64> {
self.in_flight
.lock()
.unwrap_or_else(PoisonError::into_inner)
}
fn silent(&self) -> std::sync::MutexGuard<'_, Option<SilentOpen>> {
self.silent_open
.lock()
.unwrap_or_else(PoisonError::into_inner)
}
pub fn silent_streaming(&self) -> Option<&'static str> {
self.silent().as_ref().map(|s| s.label)
}
pub fn density(&self, shape: Shape) -> f64 {
let bits = self.density[shape as usize].load(Ordering::Relaxed);
if bits == 0 {
return 1.0;
}
f64::from_bits(bits).max(1.0)
}
pub fn record_usage(&self, shape: Shape, estimate: u64, exact: u64) {
if estimate == 0 || exact == 0 {
return;
}
let ratio = (exact as f64 / estimate as f64).max(1.0);
self.density[shape as usize].store(ratio.to_bits(), Ordering::Relaxed);
}
pub fn price(&self, shape: Shape, estimate: u64, floor: u64, reply_cap: Option<u64>) -> u64 {
let prompt = ((estimate as f64 * self.density(shape)).round() as u64).max(floor);
match (self.pool, reply_cap) {
(Some(pool), None) => pool,
(_, cap) => prompt.saturating_add(cap.unwrap_or(0)),
}
}
pub async fn acquire(&self, need: u64, cancel: &CancellationToken) -> Option<Reservation<'_>> {
self.acquire_in(&self.permits, need, cancel, None, false)
.await
}
pub async fn acquire_silent(
&self,
need: u64,
cancel: &CancellationToken,
label: &'static str,
yields: bool,
) -> Option<Reservation<'_>> {
self.acquire_in(&self.silent_permits, need, cancel, Some(label), yields)
.await
}
async fn acquire_in<'a>(
&'a self,
permits: &'a Semaphore,
need: u64,
cancel: &CancellationToken,
label: Option<&'static str>,
yields: bool,
) -> Option<Reservation<'a>> {
let permit = tokio::select! {
biased;
_ = cancel.cancelled() => return None,
permit = permits.acquire() => permit.ok()?,
};
let child = label.map(|_| cancel.child_token());
let Some(pool) = self.pool else {
self.admit_silent(label, child.as_ref(), 0, yields);
return Some(Reservation {
budget: self,
tokens: 0,
label,
cancel: cancel.clone(),
child,
_permit: permit,
});
};
let mut waited = false;
let mut displacing: Option<Displacing<'_>> = None;
loop {
let notified = self.room.notified();
tokio::pin!(notified);
notified.as_mut().enable();
{
let mut open = self.open();
let fits = *open == 0 || open.saturating_add(need) <= pool;
let deferred = label.is_some() && self.displacing.load(Ordering::SeqCst) > 0;
if fits && !deferred {
*open += need;
drop(open);
self.admit_silent(label, child.as_ref(), need, yields);
drop(displacing);
return Some(Reservation {
budget: self,
tokens: need,
label,
cancel: cancel.clone(),
child,
_permit: permit,
});
}
if label.is_none()
&& !fits
&& displacing.is_none()
&& self.displace(*open, need, pool)
{
displacing = Some(Displacing::new(self));
}
if !waited {
tracing::info!(
need,
in_flight = *open,
pool,
lane = label.unwrap_or("interactive"),
"stream waits for room in the KV pool"
);
waited = true;
}
}
tokio::select! {
biased;
_ = cancel.cancelled() => return None,
_ = &mut notified => {}
}
}
}
fn admit_silent(
&self,
label: Option<&'static str>,
child: Option<&CancellationToken>,
tokens: u64,
yields: bool,
) {
if let (Some(label), Some(token)) = (label, child) {
*self.silent() = Some(SilentOpen {
label,
token: token.clone(),
tokens,
yields,
});
}
}
fn displace(&self, open: u64, need: u64, pool: u64) -> bool {
let silent = self.silent();
let Some(s) = silent.as_ref() else {
return false;
};
if !s.yields || s.tokens == 0 {
return false;
}
let without = open.saturating_sub(s.tokens);
if without != 0 && without.saturating_add(need) > pool {
return false;
}
if !s.token.is_cancelled() {
s.token.cancel();
tracing::info!(
need,
in_flight = open,
pool,
displaced = s.label,
"an interactive stream displaces the silent one"
);
}
true
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use super::*;
async fn pending<T>(fut: impl std::future::Future<Output = T>) -> bool {
tokio::time::timeout(Duration::from_millis(60), fut)
.await
.is_err()
}
#[tokio::test]
async fn a_session_is_taken_for_the_stream_and_returned_after() {
let budget = SessionBudget::new(1, None);
let cancel = CancellationToken::new();
let r = budget.acquire(10, &cancel).await;
assert!(r.is_some());
assert_eq!(budget.available_sessions(), 0);
drop(r);
assert_eq!(budget.available_sessions(), 1);
}
#[tokio::test]
async fn cancellation_ends_the_wait_for_a_session() {
let budget = SessionBudget::new(1, None);
let held = budget.acquire(0, &CancellationToken::new()).await;
let cancel = CancellationToken::new();
cancel.cancel();
assert!(budget.acquire(0, &cancel).await.is_none());
drop(held);
assert!(budget.acquire(0, &CancellationToken::new()).await.is_some());
}
#[tokio::test]
async fn reservations_that_fit_together_run_together() {
let budget = SessionBudget::new(2, Some(1000));
let cancel = CancellationToken::new();
let a = budget.acquire(400, &cancel).await.unwrap();
let b = budget.acquire(400, &cancel).await.unwrap();
assert_eq!(budget.in_flight(), 800);
drop(a);
assert_eq!(budget.in_flight(), 400);
drop(b);
assert_eq!((budget.in_flight(), budget.available_sessions()), (0, 2));
}
#[tokio::test]
async fn a_reservation_that_does_not_fit_waits_for_room() {
let budget = SessionBudget::new(3, Some(1000));
let cancel = CancellationToken::new();
let a = budget.acquire(400, &cancel).await.unwrap();
let _b = budget.acquire(400, &cancel).await.unwrap();
let mut c = std::pin::pin!(budget.acquire(400, &cancel));
assert!(pending(&mut c).await, "no room: 800 + 400 > 1000");
assert_eq!(
budget.available_sessions(),
0,
"the waiter holds its permit"
);
drop(a);
let c = c.await.expect("admitted once a sibling released");
assert_eq!(budget.in_flight(), 800);
drop(c);
}
#[tokio::test]
async fn larger_than_the_pool_is_admitted_alone_and_waits_otherwise() {
let budget = SessionBudget::new(2, Some(1000));
let cancel = CancellationToken::new();
let big = budget
.acquire(5000, &cancel)
.await
.expect("alone: admitted");
let mut small = std::pin::pin!(budget.acquire(1, &cancel));
assert!(pending(&mut small).await, "nothing fits next to 5000");
drop(big);
let small = small.await.unwrap();
let mut big2 = std::pin::pin!(budget.acquire(5000, &cancel));
assert!(pending(&mut big2).await, "not alone: waits");
drop(small);
assert!(big2.await.is_some());
}
#[tokio::test]
async fn a_cancelled_waiter_leaves_no_reservation_behind() {
let budget = SessionBudget::new(2, Some(1000));
let calm = CancellationToken::new();
let _a = budget.acquire(900, &calm).await.unwrap();
let cancel = CancellationToken::new();
let mut w = std::pin::pin!(budget.acquire(400, &cancel));
assert!(pending(&mut w).await);
cancel.cancel();
assert!(w.await.is_none());
assert_eq!(budget.in_flight(), 900);
assert_eq!(budget.available_sessions(), 1);
}
#[tokio::test]
async fn without_a_pool_only_the_count_bounds() {
let budget = SessionBudget::new(2, None);
let cancel = CancellationToken::new();
let _a = budget.acquire(u64::MAX / 2, &cancel).await.unwrap();
let _b = budget.acquire(u64::MAX / 2, &cancel).await.unwrap();
assert_eq!(budget.in_flight(), 0);
let mut c = std::pin::pin!(budget.acquire(1, &cancel));
assert!(pending(&mut c).await, "two permits, both held");
assert_eq!(SessionBudget::new(2, Some(0)).pool(), None);
}
#[tokio::test]
async fn the_count_bounds_under_a_roomy_pool() {
let budget = SessionBudget::new(1, Some(100_000));
let cancel = CancellationToken::new();
let _a = budget.acquire(1, &cancel).await.unwrap();
let mut b = std::pin::pin!(budget.acquire(1, &cancel));
assert!(pending(&mut b).await);
}
#[tokio::test]
async fn the_silent_lane_is_one_stream_wide() {
let budget = SessionBudget::new(4, None);
let cancel = CancellationToken::new();
let a = budget
.acquire_silent(10, &cancel, "reflection", false)
.await
.unwrap();
assert_eq!(budget.silent_streaming(), Some("reflection"));
let mut b = std::pin::pin!(budget.acquire_silent(10, &cancel, "compaction", false));
assert!(pending(&mut b).await, "one silent stream at a time");
assert_eq!(
budget.available_sessions(),
4,
"the interactive lane is untouched"
);
drop(a);
let b = b.await.expect("admitted once the first ended");
assert_eq!(budget.silent_streaming(), Some("compaction"));
drop(b);
assert_eq!(budget.silent_streaming(), None);
}
#[tokio::test]
async fn the_lanes_share_the_pool() {
let budget = SessionBudget::new(2, Some(1000));
let cancel = CancellationToken::new();
let silent = budget
.acquire_silent(700, &cancel, "compaction", false)
.await
.unwrap();
assert_eq!(budget.in_flight(), 700);
let mut big = std::pin::pin!(budget.acquire(400, &cancel));
assert!(pending(&mut big).await, "700 + 400 > 1000: the turn waits");
let small = budget.acquire(200, &cancel).await.expect("fits beside it");
assert_eq!(budget.in_flight(), 900);
drop(small);
drop(silent);
let big = big.await.expect("admitted once the silent round ended");
assert_eq!(budget.in_flight(), 400);
let mut quiet = std::pin::pin!(budget.acquire_silent(700, &cancel, "reflection", false));
assert!(pending(&mut quiet).await);
assert_eq!(
budget.silent_streaming(),
None,
"waiting for room is not streaming"
);
drop(big);
assert!(quiet.await.is_some());
}
#[tokio::test]
async fn a_silent_stream_alone_is_admitted_whatever_its_size() {
let budget = SessionBudget::new(1, Some(1000));
let cancel = CancellationToken::new();
let big = budget
.acquire_silent(5000, &cancel, "compaction", false)
.await
.expect("alone: admitted");
let mut turn = std::pin::pin!(budget.acquire(1, &cancel));
assert!(pending(&mut turn).await, "nothing fits beside it");
drop(big);
assert!(turn.await.is_some());
}
#[tokio::test]
async fn a_cancelled_silent_waiter_leaves_nothing_behind() {
let budget = SessionBudget::new(1, Some(1000));
let calm = CancellationToken::new();
let _turn = budget.acquire(900, &calm).await.unwrap();
let cancel = CancellationToken::new();
let mut w = std::pin::pin!(budget.acquire_silent(400, &cancel, "title", false));
assert!(pending(&mut w).await);
cancel.cancel();
assert!(w.await.is_none());
assert_eq!(budget.in_flight(), 900);
assert_eq!(budget.silent_streaming(), None);
assert!(
budget
.acquire_silent(1, &CancellationToken::new(), "title", false)
.await
.is_some(),
"the silent permit came back"
);
}
#[test]
fn pricing_calibrates_floors_and_reserves_the_pool_without_a_cap() {
let _ = Shape::COUNT;
let budget = SessionBudget::new(2, Some(1000));
assert_eq!(budget.density(Shape::Turn), 1.0);
assert_eq!(budget.price(Shape::Turn, 100, 0, Some(50)), 150);
budget.record_usage(Shape::Turn, 100, 220);
assert_eq!(budget.density(Shape::Turn), 2.2);
assert_eq!(budget.price(Shape::Turn, 100, 0, Some(50)), 270);
budget.record_usage(Shape::Turn, 100, 80);
assert_eq!(budget.density(Shape::Turn), 1.0);
assert_eq!(
budget.price(Shape::Turn, 100, 300, Some(50)),
350,
"the floor wins"
);
assert_eq!(
budget.price(Shape::Turn, 100, 0, None),
1000,
"no cap: the pool"
);
budget.record_usage(Shape::Turn, 0, 500);
budget.record_usage(Shape::Turn, 500, 0);
assert_eq!(budget.density(Shape::Turn), 1.0);
assert_eq!(
SessionBudget::new(1, None).price(Shape::Turn, 100, 0, None),
100
);
}
#[tokio::test]
async fn an_interactive_waiter_displaces_the_silent_stream_it_would_then_fit_beside() {
let budget = SessionBudget::new(1, Some(1000));
let calm = CancellationToken::new();
let roll = budget
.acquire_silent(600, &calm, "compaction", true)
.await
.unwrap();
let token = roll.stream_token();
let mut turn = std::pin::pin!(budget.acquire(600, &calm));
assert!(pending(&mut turn).await, "no room beside the roll");
assert!(token.is_cancelled(), "the roll was told to yield");
assert!(roll.displaced());
assert!(!calm.is_cancelled(), "the holder's own token is untouched");
assert_eq!(budget.displacing(), 1);
drop(roll);
let turn = turn
.await
.expect("admitted once the roll's reservation dropped");
assert_eq!(budget.displacing(), 0, "uncounted at admission");
assert_eq!(budget.in_flight(), 600);
drop(turn);
}
#[tokio::test]
async fn a_waiter_that_would_still_not_fit_displaces_nothing() {
let budget = SessionBudget::new(2, Some(1000));
let calm = CancellationToken::new();
let a = budget.acquire(500, &calm).await.unwrap();
let quiet = budget
.acquire_silent(300, &calm, "reflection", true)
.await
.unwrap();
let token = quiet.stream_token();
let mut b = std::pin::pin!(budget.acquire(600, &calm));
assert!(pending(&mut b).await, "500 + 300 + 600 > 1000");
assert!(
!token.is_cancelled(),
"500 + 600 > 1000 even without the reflection"
);
assert_eq!(budget.displacing(), 0);
drop(a);
let b = b
.await
.expect("300 + 600 fit: admitted beside the reflection");
assert!(!token.is_cancelled());
assert!(!quiet.displaced());
drop(b);
drop(quiet);
}
#[tokio::test]
async fn a_holding_silent_stream_is_never_displaced() {
let budget = SessionBudget::new(1, Some(1000));
let calm = CancellationToken::new();
let preview = budget
.acquire_silent(600, &calm, "impersonation", false)
.await
.unwrap();
let token = preview.stream_token();
let mut turn = std::pin::pin!(budget.acquire(600, &calm));
assert!(pending(&mut turn).await);
assert!(!token.is_cancelled());
assert!(!preview.displaced());
assert_eq!(budget.displacing(), 0);
drop(preview);
assert!(turn.await.is_some());
}
#[tokio::test]
async fn a_silent_waiter_displaces_nothing() {
let budget = SessionBudget::new(1, Some(1000));
let calm = CancellationToken::new();
let turn = budget.acquire(600, &calm).await.unwrap();
let mut quiet = std::pin::pin!(budget.acquire_silent(600, &calm, "title", true));
assert!(pending(&mut quiet).await);
assert!(!turn.stream_token().is_cancelled());
assert_eq!(budget.displacing(), 0);
drop(turn);
assert!(quiet.await.is_some());
}
#[tokio::test]
async fn the_displacing_waiter_is_admitted_before_the_displaced_streams_retry() {
let budget = SessionBudget::new(1, Some(1000));
let calm = CancellationToken::new();
let roll = budget
.acquire_silent(600, &calm, "compaction", true)
.await
.unwrap();
let mut turn = std::pin::pin!(budget.acquire(600, &calm));
assert!(pending(&mut turn).await);
assert!(roll.displaced());
drop(roll);
let mut retry = std::pin::pin!(budget.acquire_silent(600, &calm, "compaction", true));
assert!(
pending(&mut retry).await,
"deferred behind the displacing waiter"
);
assert_eq!(budget.in_flight(), 0);
let turn = turn.await.expect("the waiter first");
assert!(
pending(&mut retry).await,
"600 + 600 > 1000: the retry waits for the turn"
);
drop(turn);
let retry = retry.await.expect("then the retry");
assert!(!retry.displaced());
drop(retry);
}
#[tokio::test]
async fn a_displacing_waiter_that_gives_up_uncounts_itself() {
let budget = SessionBudget::new(1, Some(1000));
let calm = CancellationToken::new();
let roll = budget
.acquire_silent(600, &calm, "compaction", true)
.await
.unwrap();
let cancel = CancellationToken::new();
let mut turn = std::pin::pin!(budget.acquire(600, &cancel));
assert!(pending(&mut turn).await);
assert_eq!(budget.displacing(), 1);
cancel.cancel();
assert!(turn.await.is_none());
assert_eq!(budget.displacing(), 0);
drop(roll);
let retry = budget.acquire_silent(600, &calm, "compaction", true).await;
assert!(retry.is_some(), "nothing holds the retry back");
}
#[tokio::test]
async fn without_a_pool_nothing_is_displaced() {
let budget = SessionBudget::new(2, None);
let calm = CancellationToken::new();
let roll = budget
.acquire_silent(600, &calm, "compaction", true)
.await
.unwrap();
let turn = budget
.acquire(600, &calm)
.await
.expect("no wait without a pool");
assert!(!roll.stream_token().is_cancelled());
assert!(!roll.displaced());
assert_eq!(budget.displacing(), 0);
drop(turn);
drop(roll);
}
#[tokio::test]
async fn a_holder_that_cancelled_itself_was_not_displaced() {
let budget = SessionBudget::new(1, Some(1000));
let cancel = CancellationToken::new();
let roll = budget
.acquire_silent(600, &cancel, "compaction", true)
.await
.unwrap();
let token = roll.stream_token();
cancel.cancel();
assert!(token.is_cancelled(), "the child follows its parent");
assert!(!roll.displaced());
}
#[test]
fn a_kinds_ratio_is_its_own() {
let budget = SessionBudget::new(2, Some(100_000));
budget.record_usage(Shape::Turn, 11_058, 14_767);
budget.record_usage(Shape::Title, 294, 192);
assert!((budget.density(Shape::Turn) - 1.3354).abs() < 1e-3);
assert_eq!(budget.density(Shape::Title), 1.0, "over-counted: floored");
assert_eq!(budget.density(Shape::Roll), 1.0, "nothing recorded");
assert_eq!(budget.price(Shape::Turn, 1000, 0, Some(0)), 1335);
assert_eq!(budget.price(Shape::Title, 1000, 0, Some(0)), 1000);
assert_eq!(budget.price(Shape::Impersonation, 1000, 0, Some(0)), 1000);
}
}