use std::task::Poll;
use crate::{Error, Result, track};
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Rate(u64);
impl Rate {
pub const ZERO: Self = Self(0);
pub const fn from_bps(bps: u64) -> Self {
Self(bps)
}
pub const fn from_kbps(kbps: u64) -> Self {
Self(kbps.saturating_mul(1_000))
}
pub const fn from_mbps(mbps: u64) -> Self {
Self(mbps.saturating_mul(1_000_000))
}
pub const fn as_bps(self) -> u64 {
self.0
}
pub fn scaled(self, factor: f64) -> Self {
let scaled = self.0 as f64 * factor.max(0.0);
if scaled >= u64::MAX as f64 {
Self(u64::MAX)
} else {
Self(scaled as u64)
}
}
pub const fn abs_diff(self, other: Self) -> Self {
Self(self.0.abs_diff(other.0))
}
}
#[derive(Default)]
struct State {
bitrate: Option<Rate>,
abort: Option<Error>,
}
#[derive(Clone)]
pub struct Producer {
state: kio::Producer<State>,
}
impl Producer {
pub fn new() -> Self {
Self {
state: kio::Producer::default(),
}
}
pub fn set(&self, bitrate: Option<Rate>) -> Result<()> {
let mut state = self.modify()?;
if state.bitrate != bitrate {
state.bitrate = bitrate;
}
Ok(())
}
pub fn consume(&self) -> Consumer {
Consumer {
inner: Inner::Whole(self.state.consume()),
last: None,
}
}
pub fn abort(&self, err: Error) -> Result<()> {
let mut state = self.modify()?;
state.abort = Some(err);
state.close();
Ok(())
}
pub async fn closed(&self) -> Error {
self.state.closed().await;
self.close_error()
}
pub async fn unused(&self) -> Result<()> {
kio::wait(|waiter| self.poll_unused(waiter)).await
}
pub fn poll_unused(&self, waiter: &kio::Waiter) -> Poll<Result<()>> {
self.state.poll_unused(waiter).map(|used| match used {
Some(()) => Ok(()),
None => Err(self.close_error()),
})
}
pub fn is_used(&self) -> bool {
self.state.is_used()
}
pub async fn used(&self) -> Result<()> {
kio::wait(|waiter| self.poll_used(waiter)).await
}
pub fn poll_used(&self, waiter: &kio::Waiter) -> Poll<Result<()>> {
self.state.poll_used(waiter).map(|used| match used {
Some(()) => Ok(()),
None => Err(self.close_error()),
})
}
fn modify(&self) -> Result<kio::Mut<'_, State>> {
self.state
.write()
.map_err(|r| r.abort.clone().unwrap_or(Error::Dropped))
}
fn close_error(&self) -> Error {
self.state.read().abort.clone().unwrap_or(Error::Dropped)
}
}
impl Default for Producer {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone)]
pub struct Allocator {
estimate: Consumer,
registry: kio::Producer<Registry>,
}
impl Allocator {
pub fn new(estimate: Consumer) -> Self {
Self {
estimate,
registry: kio::Producer::default(),
}
}
pub fn unlimited() -> Self {
Self {
estimate: Consumer {
inner: Inner::Unavailable,
last: None,
},
registry: kio::Producer::default(),
}
}
pub fn reserve(&self, track: &track::Demand, max: Rate) -> Reservation {
let priority = track.priority();
let demand = track.clone();
let id = {
let Ok(mut registry) = self.registry.write() else {
unreachable!("the allocator holds its own registry producer")
};
registry.entries.retain(|entry| !entry.demand.is_closed());
let id = registry.next_id;
registry.next_id += 1;
registry.entries.push(Entry {
id,
demand,
priority,
max,
});
id
};
Reservation {
share: Share {
estimate: self.estimate.clone(),
registry: self.registry.consume(),
id,
},
registry: self.registry.downgrade(),
}
}
}
#[must_use = "a dropped Reservation is released, so the sender claims nothing and its siblings take the room"]
pub struct Reservation {
share: Share,
registry: kio::Weak<Registry>,
}
impl Reservation {
pub fn peek(&self) -> Option<Rate> {
self.share.grant()
}
pub fn consumer(&self) -> Consumer {
Consumer {
inner: Inner::Share(Box::new(self.share.clone())),
last: None,
}
}
pub fn update(&self, max: Rate) {
let Some(registry) = self.registry.upgrade() else {
return;
};
let Ok(mut registry) = registry.write() else {
return;
};
if let Some(entry) = registry.entries.iter_mut().find(|entry| entry.id == self.share.id) {
entry.max = max;
}
}
}
impl Drop for Reservation {
fn drop(&mut self) {
let Some(registry) = self.registry.upgrade() else {
return;
};
let Ok(mut registry) = registry.write() else {
return;
};
registry.entries.retain(|entry| entry.id != self.share.id);
}
}
impl Default for Allocator {
fn default() -> Self {
Self::unlimited()
}
}
impl std::fmt::Debug for Allocator {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Allocator")
.field("registered", &self.registry.read().entries.len())
.finish()
}
}
#[derive(Default)]
struct Registry {
entries: Vec<Entry>,
next_id: u64,
}
struct Entry {
id: u64,
demand: track::Demand,
priority: u8,
max: Rate,
}
#[derive(Clone)]
struct Share {
estimate: Consumer,
registry: kio::Consumer<Registry>,
id: u64,
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
struct Want {
id: u64,
priority: u8,
max: Rate,
}
impl Share {
fn grant(&self) -> Option<Rate> {
let estimate = self.estimate.peek()?;
let wants: Vec<Want> = self
.claims()
.into_iter()
.filter(|(_, demand)| demand.is_used())
.map(|(want, _)| want)
.collect();
allocate(estimate, &wants, self.id)
}
fn poll_grant(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Rate>>> {
loop {
match self.estimate.poll_changed(waiter) {
Poll::Ready(Ok(_)) => continue,
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
Poll::Pending => break,
}
}
if let Poll::Ready(Err(_)) = self.registry.poll(waiter, |_| Poll::<()>::Pending) {
return Poll::Ready(Err(Error::Dropped));
}
let wants: Vec<Want> = self
.claims()
.into_iter()
.filter(|(_, demand)| match demand.poll_state(waiter) {
track::DemandState::Active => true,
track::DemandState::Idle => false,
track::DemandState::Closed => false,
})
.map(|(want, _)| want)
.collect();
let grant = self
.estimate
.peek()
.and_then(|estimate| allocate(estimate, &wants, self.id));
Poll::Ready(Ok(grant))
}
fn claims(&self) -> Vec<(Want, track::Demand)> {
self.registry
.read()
.entries
.iter()
.map(|entry| {
(
Want {
id: entry.id,
priority: entry.priority,
max: entry.max,
},
entry.demand.clone(),
)
})
.collect()
}
}
fn allocate(estimate: Rate, wants: &[Want], id: u64) -> Option<Rate> {
let mut budget = estimate.as_bps();
let mut tier = wants.iter().map(|want| want.priority).max();
while let Some(priority) = tier {
let mut members: Vec<&Want> = wants.iter().filter(|want| want.priority == priority).collect();
members.sort_by_key(|want| want.max);
let mut remaining = members.len() as u64;
for want in members {
let even = budget / remaining;
let grant = want.max.as_bps().min(even);
if want.id == id {
return Some(Rate::from_bps(grant));
}
budget -= grant;
remaining -= 1;
}
tier = wants
.iter()
.map(|want| want.priority)
.filter(|other| *other < priority)
.max();
}
None
}
#[derive(Clone)]
pub struct Consumer {
inner: Inner,
last: Option<Rate>,
}
#[derive(Clone)]
enum Inner {
Whole(kio::Consumer<State>),
Share(Box<Share>),
Unavailable,
}
impl Consumer {
pub fn peek(&self) -> Option<Rate> {
match &self.inner {
Inner::Whole(state) => state.read().bitrate,
Inner::Share(share) => share.grant(),
Inner::Unavailable => None,
}
}
pub fn poll_changed(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Rate>>> {
let last = self.last;
let bitrate = match &mut self.inner {
Inner::Whole(state) => match state.poll(waiter, |state| {
if state.bitrate != last {
Poll::Ready(state.bitrate)
} else {
Poll::Pending
}
}) {
Poll::Ready(Ok(bitrate)) => bitrate,
Poll::Ready(Err(state)) => return Poll::Ready(Err(state.abort.clone().unwrap_or(Error::Dropped))),
Poll::Pending => return Poll::Pending,
},
Inner::Unavailable => return Poll::Pending,
Inner::Share(share) => match share.poll_grant(waiter) {
Poll::Ready(Ok(grant)) if grant == last => return Poll::Pending,
Poll::Ready(Ok(grant)) => grant,
Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
Poll::Pending => return Poll::Pending,
},
};
self.last = bitrate;
Poll::Ready(Ok(bitrate))
}
pub async fn changed(&mut self) -> Result<Option<Rate>> {
kio::wait(|waiter| self.poll_changed(waiter)).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::broadcast;
const AUDIO: u8 = 80;
const VIDEO: u8 = 60;
fn bps(bps: u64) -> Rate {
Rate::from_bps(bps)
}
fn want(id: u64, priority: u8, max: u64) -> Want {
Want {
id,
priority,
max: bps(max),
}
}
fn track(priority: u8) -> (broadcast::Producer, track::Producer) {
let broadcast = broadcast::Info::default().produce();
let track = broadcast
.create_track("t", track::Info::default().with_priority(priority))
.unwrap();
(broadcast, track)
}
#[test]
fn strict_priority_fills_the_top_tier_first() {
let wants = [want(0, AUDIO, 128_000), want(1, VIDEO, 4_000_000)];
assert_eq!(allocate(bps(2_000_000), &wants, 0), Some(bps(128_000)));
assert_eq!(allocate(bps(2_000_000), &wants, 1), Some(bps(1_872_000)));
}
#[test]
fn a_starved_tier_gets_nothing() {
let wants = [want(0, AUDIO, 2_000_000), want(1, VIDEO, 4_000_000)];
assert_eq!(allocate(bps(1_000_000), &wants, 0), Some(bps(1_000_000)));
assert_eq!(allocate(bps(1_000_000), &wants, 1), Some(bps(0)));
}
#[test]
fn one_tier_still_serves_audio_before_video() {
let flat = [want(0, 0, 128_000), want(1, 0, 4_000_000)];
let tiered = [want(0, AUDIO, 128_000), want(1, VIDEO, 4_000_000)];
for wants in [flat, tiered] {
assert_eq!(allocate(bps(2_000_000), &wants, 0), Some(bps(128_000)));
assert_eq!(allocate(bps(2_000_000), &wants, 1), Some(bps(1_872_000)));
}
}
#[test]
fn an_even_tier_splits_evenly() {
let wants = [want(0, VIDEO, 4_000_000), want(1, VIDEO, 4_000_000)];
assert_eq!(allocate(bps(6_000_000), &wants, 0), Some(bps(3_000_000)));
assert_eq!(allocate(bps(6_000_000), &wants, 1), Some(bps(3_000_000)));
}
#[test]
fn a_small_share_frees_what_it_does_not_want() {
let wants = [want(0, VIDEO, 1_000_000), want(1, VIDEO, 8_000_000)];
assert_eq!(allocate(bps(6_000_000), &wants, 0), Some(bps(1_000_000)));
assert_eq!(allocate(bps(6_000_000), &wants, 1), Some(bps(5_000_000)));
}
#[test]
fn surplus_is_left_unclaimed() {
assert_eq!(
allocate(bps(10_000_000), &[want(0, VIDEO, 4_000_000)], 0),
Some(bps(4_000_000))
);
}
#[test]
fn an_unregistered_share_has_no_grant() {
assert_eq!(allocate(bps(1_000_000), &[], 0), None);
assert_eq!(allocate(bps(1_000_000), &[want(0, VIDEO, 1_000)], 7), None);
}
#[tokio::test]
async fn concurrent_tracks_split_the_estimate() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
let (_first_broadcast, first) = track(VIDEO);
let _first_sub = first.consume();
let first = allocator.reserve(&first.demand(), bps(4_000_000));
let mut first_share = first.consumer();
estimate.set(Some(bps(2_000_000))).unwrap();
assert_eq!(first_share.changed().await.unwrap(), Some(bps(2_000_000)));
let (_second_broadcast, second) = track(VIDEO);
let _second_sub = second.consume();
let second_share = allocator.reserve(&second.demand(), bps(4_000_000));
assert_eq!(first_share.changed().await.unwrap(), Some(bps(1_000_000)));
assert_eq!(second_share.peek(), Some(bps(1_000_000)));
}
#[tokio::test]
async fn an_idle_track_claims_nothing() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
estimate.set(Some(bps(2_000_000))).unwrap();
let (_watched_broadcast, watched) = track(VIDEO);
let _watched_sub = watched.consume();
let watched_share = allocator.reserve(&watched.demand(), bps(4_000_000));
let (_idle_broadcast, idle) = track(VIDEO);
let idle_share = allocator.reserve(&idle.demand(), bps(4_000_000));
assert_eq!(watched_share.peek(), Some(bps(2_000_000)));
assert_eq!(idle_share.peek(), None);
let _idle_sub = idle.consume();
assert_eq!(watched_share.peek(), Some(bps(1_000_000)));
assert_eq!(idle_share.peek(), Some(bps(1_000_000)));
}
#[tokio::test]
async fn a_share_wakes_when_a_sibling_goes_idle() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
estimate.set(Some(bps(2_000_000))).unwrap();
let (_mine_broadcast, mine) = track(VIDEO);
let _mine_sub = mine.consume();
let mine_share_reserved = allocator.reserve(&mine.demand(), bps(4_000_000));
let mut mine_share = mine_share_reserved.consumer();
let (_sibling_broadcast, sibling) = track(VIDEO);
let sibling_sub = sibling.consume();
let _sibling_share = allocator.reserve(&sibling.demand(), bps(4_000_000));
assert_eq!(mine_share.changed().await.unwrap(), Some(bps(1_000_000)));
drop(sibling_sub);
assert_eq!(mine_share.changed().await.unwrap(), Some(bps(2_000_000)));
}
#[derive(Default)]
struct Woken(std::sync::atomic::AtomicBool);
impl std::task::Wake for Woken {
fn wake(self: std::sync::Arc<Self>) {
self.wake_by_ref();
}
fn wake_by_ref(self: &std::sync::Arc<Self>) {
self.0.store(true, std::sync::atomic::Ordering::SeqCst);
}
}
impl Woken {
fn new() -> (std::sync::Arc<Self>, kio::Waiter) {
let flag = std::sync::Arc::new(Self::default());
let waiter = kio::Waiter::new(std::task::Waker::from(flag.clone()));
(flag, waiter)
}
fn woken(&self) -> bool {
self.0.load(std::sync::atomic::Ordering::SeqCst)
}
}
#[tokio::test]
async fn an_unchanged_slice_keeps_watching_the_estimate() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
let (_broadcast, track) = track(VIDEO);
let _sub = track.consume();
let share_reserved = allocator.reserve(&track.demand(), bps(4_000_000));
let mut share = share_reserved.consumer();
estimate.set(Some(bps(10_000_000))).unwrap();
assert_eq!(share.changed().await.unwrap(), Some(bps(4_000_000)));
let (woken, waiter) = Woken::new();
estimate.set(Some(bps(9_000_000))).unwrap();
assert!(share.poll_changed(&waiter).is_pending());
estimate.set(Some(bps(1_000_000))).unwrap();
assert!(woken.woken(), "a parked share must be woken by the next estimate");
assert_eq!(share.changed().await.unwrap(), Some(bps(1_000_000)));
}
#[tokio::test]
async fn a_parked_share_is_woken_by_sibling_demand() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
estimate.set(Some(bps(2_000_000))).unwrap();
let (_mine_broadcast, mine) = track(VIDEO);
let _mine_sub = mine.consume();
let share_reserved = allocator.reserve(&mine.demand(), bps(4_000_000));
let mut share = share_reserved.consumer();
let (_sibling_broadcast, sibling) = track(VIDEO);
let sibling_sub = sibling.consume();
let _sibling_share = allocator.reserve(&sibling.demand(), bps(4_000_000));
assert_eq!(share.changed().await.unwrap(), Some(bps(1_000_000)));
let (woken, waiter) = Woken::new();
assert!(share.poll_changed(&waiter).is_pending());
drop(sibling_sub);
assert!(woken.woken(), "a sibling going idle must wake a parked share");
}
#[tokio::test]
async fn a_share_follows_the_estimate_lifecycle() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
let (_broadcast, track) = track(VIDEO);
let _sub = track.consume();
let share_reserved = allocator.reserve(&track.demand(), bps(4_000_000));
let mut share = share_reserved.consumer();
estimate.set(Some(bps(2_000_000))).unwrap();
assert_eq!(share.changed().await.unwrap(), Some(bps(2_000_000)));
estimate.set(None).unwrap();
assert_eq!(share.changed().await.unwrap(), None);
estimate.abort(Error::Cancel).unwrap();
assert!(share.changed().await.is_err());
assert!(share.changed().await.is_err());
}
#[tokio::test]
async fn a_closed_track_is_pruned() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
let (_first_broadcast, first) = track(VIDEO);
let _first_share = allocator.reserve(&first.demand(), bps(4_000_000));
first.abort(Error::Cancel).unwrap();
let (_second_broadcast, second) = track(VIDEO);
let _second_share = allocator.reserve(&second.demand(), bps(4_000_000));
assert_eq!(allocator.registry.read().entries.len(), 1);
}
#[tokio::test]
async fn dropping_a_reservation_releases_it() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
estimate.set(Some(bps(2_000_000))).unwrap();
let (_first_broadcast, first) = track(VIDEO);
let _first_sub = first.consume();
let first_reserved = allocator.reserve(&first.demand(), bps(4_000_000));
let (_second_broadcast, second) = track(VIDEO);
let _second_sub = second.consume();
let second_reserved = allocator.reserve(&second.demand(), bps(4_000_000));
assert_eq!(second_reserved.peek(), Some(bps(1_000_000)));
let mut orphan = first_reserved.consumer();
assert_eq!(orphan.changed().await.unwrap(), Some(bps(1_000_000)));
drop(first_reserved);
assert_eq!(allocator.registry.read().entries.len(), 1);
assert_eq!(second_reserved.peek(), Some(bps(2_000_000)));
assert_eq!(orphan.changed().await.unwrap(), None);
assert_eq!(orphan.peek(), None);
}
#[tokio::test]
async fn update_changes_the_claim_in_place() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
estimate.set(Some(bps(6_000_000))).unwrap();
let (_small_broadcast, small) = track(VIDEO);
let _small_sub = small.consume();
let small_reserved = allocator.reserve(&small.demand(), bps(1_000_000));
let (_large_broadcast, large) = track(VIDEO);
let _large_sub = large.consume();
let large_reserved = allocator.reserve(&large.demand(), bps(8_000_000));
assert_eq!(small_reserved.peek(), Some(bps(1_000_000)));
assert_eq!(large_reserved.peek(), Some(bps(5_000_000)));
small_reserved.update(bps(4_000_000));
assert_eq!(allocator.registry.read().entries.len(), 2);
assert_eq!(small_reserved.peek(), Some(bps(3_000_000)));
assert_eq!(large_reserved.peek(), Some(bps(3_000_000)));
small_reserved.update(bps(1_000_000));
assert_eq!(large_reserved.peek(), Some(bps(5_000_000)));
}
#[tokio::test]
async fn update_wakes_a_parked_reader() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
estimate.set(Some(bps(6_000_000))).unwrap();
let (_broadcast, producer) = track(VIDEO);
let _sub = producer.consume();
let reserved = allocator.reserve(&producer.demand(), bps(1_000_000));
let mut share = reserved.consumer();
assert_eq!(share.changed().await.unwrap(), Some(bps(1_000_000)));
let (woken, waiter) = Woken::new();
assert!(share.poll_changed(&waiter).is_pending());
reserved.update(bps(4_000_000));
assert!(woken.woken(), "raising the ceiling must wake the reader");
assert_eq!(share.changed().await.unwrap(), Some(bps(4_000_000)));
}
#[tokio::test]
async fn a_share_outliving_the_allocator_reports_closed() {
let estimate = Producer::new();
let allocator = Allocator::new(estimate.consume());
let (_broadcast, producer) = track(VIDEO);
let _sub = producer.consume();
let share_reserved = allocator.reserve(&producer.demand(), bps(4_000_000));
let mut share = share_reserved.consumer();
drop(allocator);
drop(estimate);
assert!(share.changed().await.is_err());
}
#[tokio::test]
async fn closed_is_distinct_from_unavailable() {
let producer = Producer::new();
let mut consumer = producer.consume();
producer.set(Some(bps(1_000_000))).unwrap();
assert_eq!(consumer.changed().await.unwrap(), Some(bps(1_000_000)));
producer.set(None).unwrap();
assert_eq!(consumer.changed().await.unwrap(), None);
producer.abort(Error::Cancel).unwrap();
assert!(consumer.changed().await.is_err());
assert!(consumer.changed().await.is_err());
}
}