Skip to main content

moq_net/model/
bandwidth.rs

1//! Rate estimation, split into a [Producer] and [Consumer] handle.
2//!
3//! A [Producer] is used to set the current estimated bitrate, notifying consumers.
4//! A [Consumer] can read the current estimate and wait for changes.
5//!
6//! One estimate covers a whole connection, so senders sharing one divide it with
7//! an [Allocator] rather than each targeting the whole thing. How a sender then
8//! *follows* its share is shared media policy: see `moq_mux::rate`.
9
10use std::task::Poll;
11
12use crate::{Error, Result, track};
13
14/// A rate, in bits per second.
15///
16/// A newtype rather than a bare integer because everything that meets in an
17/// [`Allocator`] is the same quantity measured the same way: a congestion
18/// controller's estimate, a track's reservation, an encoder's ceiling. One of them
19/// reading bytes per second, or kilobits, is off by a factor of eight or a thousand
20/// and still typechecks, which is the kind of wrong that reaches production.
21///
22/// Named for what it measures rather than for the module: `Session::stats` has called
23/// this quantity a rate all along (`estimated_send_rate`).
24#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
25pub struct Rate(u64);
26
27impl Rate {
28	/// No bandwidth at all.
29	pub const ZERO: Self = Self(0);
30
31	/// A rate in bits per second, the unit every wire and codec API here uses.
32	pub const fn from_bps(bps: u64) -> Self {
33		Self(bps)
34	}
35
36	/// A rate in kilobits (1000 bits) per second, saturating.
37	pub const fn from_kbps(kbps: u64) -> Self {
38		Self(kbps.saturating_mul(1_000))
39	}
40
41	/// A rate in megabits (1000 kilobits) per second, saturating.
42	pub const fn from_mbps(mbps: u64) -> Self {
43		Self(mbps.saturating_mul(1_000_000))
44	}
45
46	/// This rate in bits per second, for handing to a codec or FFI that wants a plain integer.
47	pub const fn as_bps(self) -> u64 {
48		self.0
49	}
50
51	/// This rate scaled by `factor`, saturating at both ends.
52	///
53	/// For the fractional arithmetic rate control does (a ramp allowance, a hysteresis
54	/// band) without spreading `as` casts across the callers.
55	pub fn scaled(self, factor: f64) -> Self {
56		let scaled = self.0 as f64 * factor.max(0.0);
57		if scaled >= u64::MAX as f64 {
58			Self(u64::MAX)
59		} else {
60			Self(scaled as u64)
61		}
62	}
63
64	/// The absolute difference between two rates.
65	pub const fn abs_diff(self, other: Self) -> Self {
66		Self(self.0.abs_diff(other.0))
67	}
68}
69
70#[derive(Default)]
71struct State {
72	bitrate: Option<Rate>,
73	abort: Option<Error>,
74}
75
76/// Produces bandwidth estimates, notifying consumers when the value changes.
77#[derive(Clone)]
78pub struct Producer {
79	state: kio::Producer<State>,
80}
81
82impl Producer {
83	/// Create a fresh producer with no current estimate.
84	pub fn new() -> Self {
85		Self {
86			state: kio::Producer::default(),
87		}
88	}
89
90	/// Set the current bandwidth estimate, or `None` while the backend has none.
91	pub fn set(&self, bitrate: Option<Rate>) -> Result<()> {
92		let mut state = self.modify()?;
93		if state.bitrate != bitrate {
94			state.bitrate = bitrate;
95		}
96		Ok(())
97	}
98
99	/// Create a new consumer for the bandwidth estimate.
100	pub fn consume(&self) -> Consumer {
101		Consumer {
102			inner: Inner::Whole(self.state.consume()),
103			last: None,
104		}
105	}
106
107	/// Close the producer with an error, notifying all consumers.
108	pub fn abort(&self, err: Error) -> Result<()> {
109		let mut state = self.modify()?;
110		state.abort = Some(err);
111		state.close();
112		Ok(())
113	}
114
115	/// Block until the channel is closed, returning the cause.
116	pub async fn closed(&self) -> Error {
117		self.state.closed().await;
118		self.close_error()
119	}
120
121	/// Block until there are no active consumers.
122	pub async fn unused(&self) -> Result<()> {
123		kio::wait(|waiter| self.poll_unused(waiter)).await
124	}
125
126	/// Poll until there are no active consumers. Errors if the channel closes first.
127	pub fn poll_unused(&self, waiter: &kio::Waiter) -> Poll<Result<()>> {
128		self.state.poll_unused(waiter).map(|used| match used {
129			Some(()) => Ok(()),
130			None => Err(self.close_error()),
131		})
132	}
133
134	/// Whether at least one active consumer exists right now.
135	pub fn is_used(&self) -> bool {
136		self.state.is_used()
137	}
138
139	/// Block until there is at least one active consumer.
140	pub async fn used(&self) -> Result<()> {
141		kio::wait(|waiter| self.poll_used(waiter)).await
142	}
143
144	/// Poll until at least one active consumer exists. Errors if the channel closes first.
145	pub fn poll_used(&self, waiter: &kio::Waiter) -> Poll<Result<()>> {
146		self.state.poll_used(waiter).map(|used| match used {
147			Some(()) => Ok(()),
148			None => Err(self.close_error()),
149		})
150	}
151
152	fn modify(&self) -> Result<kio::Mut<'_, State>> {
153		self.state
154			.write()
155			.map_err(|r| r.abort.clone().unwrap_or(Error::Dropped))
156	}
157
158	/// The close error, once the channel is closed.
159	fn close_error(&self) -> Error {
160		self.state.read().abort.clone().unwrap_or(Error::Dropped)
161	}
162}
163
164impl Default for Producer {
165	fn default() -> Self {
166		Self::new()
167	}
168}
169
170/// Divides one connection's bandwidth estimate among the tracks sharing it.
171///
172/// Every sender on a connection reads the same estimate, so N senders each
173/// targeting all of it oversubscribe the uplink N times over. Register a track
174/// here and it gets a [`Consumer`] reporting only its own slice, so the slices
175/// sum to the estimate instead of each matching it.
176///
177/// Advisory, not enforced. A track that ignores its slice, or can't follow it at
178/// all (PCM audio has a fixed bitrate), still sends what it sends; the transport
179/// sheds the excess by dropping groups. Rate estimation isn't exact enough
180/// for the difference to be worth policing.
181///
182/// Clones share one registry, so hand a clone to each sender.
183#[derive(Clone)]
184pub struct Allocator {
185	estimate: Consumer,
186	registry: kio::Producer<Registry>,
187}
188
189impl Allocator {
190	/// Divide `estimate`, normally a connection's
191	/// [`Session::send_bandwidth`](crate::Session::send_bandwidth).
192	pub fn new(estimate: Consumer) -> Self {
193		Self {
194			estimate,
195			registry: kio::Producer::default(),
196		}
197	}
198
199	/// An allocator with nothing to divide, so every reservation reports `None`.
200	///
201	/// `None` already means "no opinion, hold your rate" to a sender, so this is what
202	/// a transport with no congestion estimate, a local file, or a test harness wants,
203	/// and it saves every config struct on the way down from being an `Option`. It is
204	/// also [`Default`], so a config that never sets one encodes at its configured rate.
205	pub fn unlimited() -> Self {
206		Self {
207			estimate: Consumer {
208				inner: Inner::Unavailable,
209				last: None,
210			},
211			registry: kio::Producer::default(),
212		}
213	}
214
215	/// Reserve up to `max` for `track`, returning the reservation.
216	///
217	/// `max` is a ceiling, not a measurement: reserve the most the track can ever
218	/// send, not what it happens to be sending. A VBR encoder sitting on a black
219	/// screen at 1 Mbps can jump to 6 Mbps between one frame and the next, and a
220	/// reservation that had followed it down would have already handed that room
221	/// to somebody else.
222	///
223	/// Priority comes from the track ([`track::Info::priority`], higher served
224	/// first). A tier is filled to its reservations before the next one sees a
225	/// bit; within a tier the split is max-min fair, so a share asking for less
226	/// than an even cut takes all of it and leaves the difference to the others.
227	///
228	/// That is the *publisher's* priority, which is not what orders the local send
229	/// queue: that ranks by each subscription's own priority, so a subscriber
230	/// asking for video ahead of audio is served that way whatever this decides.
231	/// The publisher's is still the right one to divide by, since allocation is a
232	/// decision about what to *produce*, and there is no single subscriber
233	/// priority to read when several are watching one track.
234	///
235	/// That last part is what carries the common case, since publishers leave
236	/// `priority` at its default today: one tier of audio and video still serves
237	/// audio's small reservation in full before video takes the remainder.
238	///
239	/// The reservation lasts as long as the returned [`Reservation`]: hold it for as
240	/// long as the sender is publishing, change the ceiling with
241	/// [`update`](Reservation::update), and drop it to hand the room back. It is
242	/// released when the track closes either way, since a [`track::Demand`] is a weak
243	/// handle and reserving never keeps a track alive.
244	///
245	/// Read the current slice through [`Reservation::consumer`], which reports `None`
246	/// while nothing is subscribed to the track or the connection has no estimate.
247	/// That tells a sender to hold its current rate rather than encode at zero.
248	pub fn reserve(&self, track: &track::Demand, max: Rate) -> Reservation {
249		// Read the track before taking the registry lock, so the two are never held
250		// at once and there's no order to get wrong.
251		let priority = track.priority();
252		let demand = track.clone();
253
254		let id = {
255			// Nothing ever closes this channel: the allocator holds the only producer
256			// and never aborts it, so it's open for as long as `self` is.
257			let Ok(mut registry) = self.registry.write() else {
258				unreachable!("the allocator holds its own registry producer")
259			};
260			// Closed tracks can't be demanded again; drop them rather than walking
261			// them on every poll for the rest of the connection.
262			registry.entries.retain(|entry| !entry.demand.is_closed());
263
264			let id = registry.next_id;
265			registry.next_id += 1;
266			registry.entries.push(Entry {
267				id,
268				demand,
269				priority,
270				max,
271			});
272			id
273		};
274
275		Reservation {
276			share: Share {
277				estimate: self.estimate.clone(),
278				registry: self.registry.consume(),
279				id,
280			},
281			registry: self.registry.downgrade(),
282		}
283	}
284}
285
286/// One track's standing claim on an [`Allocator`], held for as long as the sender
287/// that took it is publishing.
288///
289/// Separate from the [`Consumer`] that reads the slice, because the two have opposite
290/// lifetimes: read handles are cloned around and dropped freely, while the claim itself
291/// has to outlive every one of them or the sender silently stops claiming anything.
292#[must_use = "a dropped Reservation is released, so the sender claims nothing and its siblings take the room"]
293pub struct Reservation {
294	share: Share,
295	/// Weak so a reservation can't keep the registry alive: one outliving every
296	/// [`Allocator`] reports the estimate as gone rather than holding the channel open.
297	registry: kio::Weak<Registry>,
298}
299
300impl Reservation {
301	/// This reservation's slice right now.
302	///
303	/// Stateless, unlike [`Consumer::changed`], which carries a cursor over what it last
304	/// reported and so needs a handle of its own.
305	pub fn peek(&self) -> Option<Rate> {
306		self.share.grant()
307	}
308
309	/// A handle reading this reservation's current slice of the estimate.
310	///
311	/// Cloneable and independent of the reservation: once the reservation is dropped
312	/// these report `None`, the same "hold your rate" a track nobody is watching reports.
313	pub fn consumer(&self) -> Consumer {
314		Consumer {
315			inner: Inner::Share(Box::new(self.share.clone())),
316			last: None,
317		}
318	}
319
320	/// Change the ceiling, keeping the same claim.
321	///
322	/// For a sender whose ceiling genuinely moved: an encoder reopening at a resolution
323	/// it negotiated with the device, not an encoder observing its own output. A
324	/// reservation that followed the rate a VBR source happens to be sending would hand
325	/// the room away every time the picture went still, and not have it back when the
326	/// picture moved again.
327	///
328	/// Does nothing once every [`Allocator`] is gone, since there is then nothing left
329	/// dividing anything; a reader learns that from its [`consumer`](Self::consumer).
330	pub fn update(&self, max: Rate) {
331		let Some(registry) = self.registry.upgrade() else {
332			return;
333		};
334		let Ok(mut registry) = registry.write() else {
335			return;
336		};
337		if let Some(entry) = registry.entries.iter_mut().find(|entry| entry.id == self.share.id) {
338			entry.max = max;
339		}
340	}
341}
342
343impl Drop for Reservation {
344	fn drop(&mut self) {
345		let Some(registry) = self.registry.upgrade() else {
346			return;
347		};
348		let Ok(mut registry) = registry.write() else {
349			return;
350		};
351		registry.entries.retain(|entry| entry.id != self.share.id);
352	}
353}
354
355/// An allocator with nothing to divide, so a config that never sets one leaves its
356/// senders at their configured rates. See [`Allocator::unlimited`].
357impl Default for Allocator {
358	fn default() -> Self {
359		Self::unlimited()
360	}
361}
362
363// Hand-written so the config structs that carry an allocator can still derive
364// `Debug`. The registry is the only part worth printing.
365impl std::fmt::Debug for Allocator {
366	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
367		f.debug_struct("Allocator")
368			.field("registered", &self.registry.read().entries.len())
369			.finish()
370	}
371}
372
373/// Every track registered with one [`Allocator`].
374#[derive(Default)]
375struct Registry {
376	entries: Vec<Entry>,
377	next_id: u64,
378}
379
380/// One registered track's standing reservation.
381struct Entry {
382	id: u64,
383	demand: track::Demand,
384	priority: u8,
385	max: Rate,
386}
387
388/// A share's view of the estimate it divides.
389#[derive(Clone)]
390struct Share {
391	/// What's being divided, which may itself be a share.
392	estimate: Consumer,
393	registry: kio::Consumer<Registry>,
394	id: u64,
395}
396
397/// One demanded track's claim, snapshotted out of the [`Registry`] so no lock is
398/// held while the tracks themselves are read.
399#[derive(Copy, Clone, Debug, Eq, PartialEq)]
400struct Want {
401	id: u64,
402	priority: u8,
403	max: Rate,
404}
405
406impl Share {
407	/// This share's slice right now.
408	fn grant(&self) -> Option<Rate> {
409		let estimate = self.estimate.peek()?;
410		let wants: Vec<Want> = self
411			.claims()
412			.into_iter()
413			.filter(|(_, demand)| demand.is_used())
414			.map(|(want, _)| want)
415			.collect();
416		allocate(estimate, &wants, self.id)
417	}
418
419	/// This share's slice, arming `waiter` for everything that could move it: the
420	/// estimate, the set of registered tracks, and each track's demand.
421	fn poll_grant(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Rate>>> {
422		// Drain the estimate rather than polling it once. A poll that returns Ready
423		// registers no waker, and this share may still conclude its own slice didn't
424		// move (its reservation caps it, so most estimate changes don't reach it) and
425		// park. Without the drain that park would never be woken by the next move.
426		// The value itself is read below, since an unchanged estimate still needs
427		// re-dividing when the tracks sharing it change.
428		loop {
429			match self.estimate.poll_changed(waiter) {
430				// Moved: go round again, which either finds it settled and arms the
431				// waker, or finds it moved again and makes progress toward the latest.
432				Poll::Ready(Ok(_)) => continue,
433				Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
434				Poll::Pending => break,
435			}
436		}
437
438		// Wake on any registration change. The closure never completes, so the only
439		// ready outcome is the registry being gone, i.e. every allocator dropped.
440		if let Poll::Ready(Err(_)) = self.registry.poll(waiter, |_| Poll::<()>::Pending) {
441			return Poll::Ready(Err(Error::Dropped));
442		}
443
444		let wants: Vec<Want> = self
445			.claims()
446			.into_iter()
447			.filter(|(_, demand)| match demand.poll_state(waiter) {
448				track::DemandState::Active => true,
449				// An idle track claims nothing, so an unwatched encoder's reservation
450				// goes to whoever is actually sending.
451				track::DemandState::Idle => false,
452				// Gone for good, and pruned by the next `register`.
453				track::DemandState::Closed => false,
454			})
455			.map(|(want, _)| want)
456			.collect();
457
458		let grant = self
459			.estimate
460			.peek()
461			.and_then(|estimate| allocate(estimate, &wants, self.id));
462		Poll::Ready(Ok(grant))
463	}
464
465	/// Snapshot the registry so the tracks can be read without holding its lock.
466	fn claims(&self) -> Vec<(Want, track::Demand)> {
467		self.registry
468			.read()
469			.entries
470			.iter()
471			.map(|entry| {
472				(
473					Want {
474						id: entry.id,
475						priority: entry.priority,
476						max: entry.max,
477					},
478					entry.demand.clone(),
479				)
480			})
481			.collect()
482	}
483}
484
485/// Divide `estimate` among `wants`, returning the slice for `id`.
486///
487/// Strict priority: a tier is filled to its reservations before the next tier
488/// sees a bit. Within a tier the split is max-min fair, so a share asking for
489/// less than an even split takes all of it and leaves the rest to the others.
490///
491/// Surplus above the total reserved is left unclaimed rather than spread around.
492/// A reservation is what a sender can use, so handing it more is not a reason to
493/// send more than it was configured for.
494///
495/// `None` when `id` isn't among the wants, which is how an idle or closed track
496/// reports "hold your rate" instead of a grant of zero.
497fn allocate(estimate: Rate, wants: &[Want], id: u64) -> Option<Rate> {
498	// Plain bits per second inside: this is where the division actually happens, and
499	// wrapping every intermediate would need arithmetic on `Rate` that no caller wants.
500	let mut budget = estimate.as_bps();
501	let mut tier = wants.iter().map(|want| want.priority).max();
502
503	while let Some(priority) = tier {
504		// Ascending by reservation: each share takes an even cut of what's left, or
505		// all it asked for if that's less, which frees the difference for the rest.
506		let mut members: Vec<&Want> = wants.iter().filter(|want| want.priority == priority).collect();
507		members.sort_by_key(|want| want.max);
508
509		let mut remaining = members.len() as u64;
510		for want in members {
511			let even = budget / remaining;
512			let grant = want.max.as_bps().min(even);
513			if want.id == id {
514				return Some(Rate::from_bps(grant));
515			}
516			budget -= grant;
517			remaining -= 1;
518		}
519
520		tier = wants
521			.iter()
522			.map(|want| want.priority)
523			.filter(|other| *other < priority)
524			.max();
525	}
526
527	None
528}
529
530/// Consumes bandwidth estimates, allowing reads and async change notifications.
531#[derive(Clone)]
532pub struct Consumer {
533	inner: Inner,
534	last: Option<Rate>,
535}
536
537/// What a [`Consumer`] is reading: the whole estimate, or one track's slice of it.
538#[derive(Clone)]
539enum Inner {
540	Whole(kio::Consumer<State>),
541	Share(Box<Share>),
542	/// [`Allocator::unlimited`]'s: no estimate, and never will be.
543	Unavailable,
544}
545
546impl Consumer {
547	/// Get the current bandwidth estimate synchronously.
548	pub fn peek(&self) -> Option<Rate> {
549		match &self.inner {
550			Inner::Whole(state) => state.read().bitrate,
551			Inner::Share(share) => share.grant(),
552			Inner::Unavailable => None,
553		}
554	}
555
556	/// Poll for a bandwidth change without blocking.
557	///
558	/// `Ok(None)` means the estimate is unavailable *for now*: the backend
559	/// stopped reporting one, or the handle spans reconnects and is between
560	/// sessions. `Err` means the producer is gone and no further change will ever
561	/// arrive. They're distinct because a caller holds its current rate for the
562	/// first and stops watching for the second.
563	///
564	/// A backend with no bandwidth estimation at all yields no [Consumer] in the
565	/// first place, so that case never reaches here.
566	pub fn poll_changed(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Rate>>> {
567		let last = self.last;
568
569		let bitrate = match &mut self.inner {
570			Inner::Whole(state) => match state.poll(waiter, |state| {
571				if state.bitrate != last {
572					Poll::Ready(state.bitrate)
573				} else {
574					Poll::Pending
575				}
576			}) {
577				Poll::Ready(Ok(bitrate)) => bitrate,
578				// Closed, and the value hasn't moved since the last read: report it as
579				// terminal. Collapsing this into `Ok(None)` would be indistinguishable
580				// from a live-but-unavailable estimate, and since a closed channel is
581				// always immediately ready, a `select!` over it would spin forever.
582				Poll::Ready(Err(state)) => return Poll::Ready(Err(state.abort.clone().unwrap_or(Error::Dropped))),
583				Poll::Pending => return Poll::Pending,
584			},
585			// A share recomputes its slice on every wakeup, so it filters the
586			// unchanged case here rather than inside the poll. Every waker that could
587			// move the slice was armed by `poll_grant` either way.
588			// Nothing to report and nothing that could ever change it, so park without
589			// arming anything rather than reporting a `None` the caller would re-read forever.
590			Inner::Unavailable => return Poll::Pending,
591			Inner::Share(share) => match share.poll_grant(waiter) {
592				Poll::Ready(Ok(grant)) if grant == last => return Poll::Pending,
593				Poll::Ready(Ok(grant)) => grant,
594				Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
595				Poll::Pending => return Poll::Pending,
596			},
597		};
598
599		self.last = bitrate;
600		Poll::Ready(Ok(bitrate))
601	}
602
603	/// Block until the bandwidth estimate changes, returning the new value, or
604	/// `None` when the estimate has become unavailable.
605	///
606	/// # Errors
607	///
608	/// Returns an error once the producer is closed or dropped, so a caller can
609	/// stop watching. See [`poll_changed`](Self::poll_changed).
610	pub async fn changed(&mut self) -> Result<Option<Rate>> {
611		kio::wait(|waiter| self.poll_changed(waiter)).await
612	}
613}
614
615#[cfg(test)]
616mod tests {
617	use super::*;
618	use crate::broadcast;
619
620	/// Priorities matching `hang`'s, which is what the allocator sees in practice.
621	const AUDIO: u8 = 80;
622	const VIDEO: u8 = 60;
623
624	/// Bits per second, so the tables below stay readable.
625	fn bps(bps: u64) -> Rate {
626		Rate::from_bps(bps)
627	}
628
629	fn want(id: u64, priority: u8, max: u64) -> Want {
630		Want {
631			id,
632			priority,
633			max: bps(max),
634		}
635	}
636
637	/// A standalone track at `priority`, plus the broadcast keeping it alive.
638	fn track(priority: u8) -> (broadcast::Producer, track::Producer) {
639		let broadcast = broadcast::Info::default().produce();
640		let track = broadcast
641			.create_track("t", track::Info::default().with_priority(priority))
642			.unwrap();
643		(broadcast, track)
644	}
645
646	#[test]
647	fn strict_priority_fills_the_top_tier_first() {
648		let wants = [want(0, AUDIO, 128_000), want(1, VIDEO, 4_000_000)];
649
650		// Audio takes its reservation off the top; video gets what's left. This is
651		// the job `rate::Policy::headroom` used to approximate with a flat 10%.
652		assert_eq!(allocate(bps(2_000_000), &wants, 0), Some(bps(128_000)));
653		assert_eq!(allocate(bps(2_000_000), &wants, 1), Some(bps(1_872_000)));
654	}
655
656	#[test]
657	fn a_starved_tier_gets_nothing() {
658		let wants = [want(0, AUDIO, 2_000_000), want(1, VIDEO, 4_000_000)];
659
660		assert_eq!(allocate(bps(1_000_000), &wants, 0), Some(bps(1_000_000)));
661		// Strict, not weighted: the lower tier is not owed a floor. Its encoder
662		// clamps to `rate::Policy::min` and the transport sheds what won't fit.
663		assert_eq!(allocate(bps(1_000_000), &wants, 1), Some(bps(0)));
664	}
665
666	/// Publishers don't set [`track::Info::priority`] today (it defaults to 0 and
667	/// `hang::container::track_info` leaves it there), so audio and video land in
668	/// one tier. That has to come out right anyway, and it does: max-min fair
669	/// satisfies the small claim first, so audio still gets its full reservation
670	/// and video takes the rest. Priority only changes the answer once a tier's
671	/// smaller claims outgrow an even split.
672	#[test]
673	fn one_tier_still_serves_audio_before_video() {
674		let flat = [want(0, 0, 128_000), want(1, 0, 4_000_000)];
675		let tiered = [want(0, AUDIO, 128_000), want(1, VIDEO, 4_000_000)];
676
677		for wants in [flat, tiered] {
678			assert_eq!(allocate(bps(2_000_000), &wants, 0), Some(bps(128_000)));
679			assert_eq!(allocate(bps(2_000_000), &wants, 1), Some(bps(1_872_000)));
680		}
681	}
682
683	#[test]
684	fn an_even_tier_splits_evenly() {
685		let wants = [want(0, VIDEO, 4_000_000), want(1, VIDEO, 4_000_000)];
686
687		assert_eq!(allocate(bps(6_000_000), &wants, 0), Some(bps(3_000_000)));
688		assert_eq!(allocate(bps(6_000_000), &wants, 1), Some(bps(3_000_000)));
689	}
690
691	/// Max-min fair, not an even split: a 360p rung sharing with a 1080p one takes
692	/// only what it asked for and leaves the rest, instead of both being held to half.
693	#[test]
694	fn a_small_share_frees_what_it_does_not_want() {
695		let wants = [want(0, VIDEO, 1_000_000), want(1, VIDEO, 8_000_000)];
696
697		assert_eq!(allocate(bps(6_000_000), &wants, 0), Some(bps(1_000_000)));
698		assert_eq!(allocate(bps(6_000_000), &wants, 1), Some(bps(5_000_000)));
699	}
700
701	/// Capping at the reservation is what keeps an uncongested link encoding at
702	/// exactly the configured rate. Spreading the surplus instead would only matter
703	/// if senders scaled a fraction of their grant, which is what headroom did.
704	#[test]
705	fn surplus_is_left_unclaimed() {
706		assert_eq!(
707			allocate(bps(10_000_000), &[want(0, VIDEO, 4_000_000)], 0),
708			Some(bps(4_000_000))
709		);
710	}
711
712	#[test]
713	fn an_unregistered_share_has_no_grant() {
714		assert_eq!(allocate(bps(1_000_000), &[], 0), None);
715		assert_eq!(allocate(bps(1_000_000), &[want(0, VIDEO, 1_000)], 7), None);
716	}
717
718	/// The headline case: two encoders on one connection must not each target the
719	/// whole estimate.
720	#[tokio::test]
721	async fn concurrent_tracks_split_the_estimate() {
722		let estimate = Producer::new();
723		let allocator = Allocator::new(estimate.consume());
724
725		let (_first_broadcast, first) = track(VIDEO);
726		let _first_sub = first.consume();
727		let first = allocator.reserve(&first.demand(), bps(4_000_000));
728		let mut first_share = first.consumer();
729
730		estimate.set(Some(bps(2_000_000))).unwrap();
731		// Alone, it gets everything it asked for that the link can carry.
732		assert_eq!(first_share.changed().await.unwrap(), Some(bps(2_000_000)));
733
734		// A second encoder starts. The first must give half back without anyone
735		// telling it to, or the two together target 200% of the uplink.
736		let (_second_broadcast, second) = track(VIDEO);
737		let _second_sub = second.consume();
738		let second_share = allocator.reserve(&second.demand(), bps(4_000_000));
739
740		assert_eq!(first_share.changed().await.unwrap(), Some(bps(1_000_000)));
741		assert_eq!(second_share.peek(), Some(bps(1_000_000)));
742	}
743
744	/// An unwatched track releases its reservation, since `publish_capture` stops
745	/// encoding entirely while nothing is subscribed.
746	#[tokio::test]
747	async fn an_idle_track_claims_nothing() {
748		let estimate = Producer::new();
749		let allocator = Allocator::new(estimate.consume());
750		estimate.set(Some(bps(2_000_000))).unwrap();
751
752		let (_watched_broadcast, watched) = track(VIDEO);
753		let _watched_sub = watched.consume();
754		let watched_share = allocator.reserve(&watched.demand(), bps(4_000_000));
755
756		let (_idle_broadcast, idle) = track(VIDEO);
757		let idle_share = allocator.reserve(&idle.demand(), bps(4_000_000));
758
759		assert_eq!(watched_share.peek(), Some(bps(2_000_000)));
760		// Not zero: an idle share reports "no opinion" so a sender that is mid-shutdown
761		// holds its rate instead of retuning to the floor on the way out.
762		assert_eq!(idle_share.peek(), None);
763
764		// It joins, and the split happens.
765		let _idle_sub = idle.consume();
766		assert_eq!(watched_share.peek(), Some(bps(1_000_000)));
767		assert_eq!(idle_share.peek(), Some(bps(1_000_000)));
768	}
769
770	/// Demand transitions have to wake a parked share, not just change what a later
771	/// `peek` would see; the encoder is sitting in a `select!` on `changed`.
772	#[tokio::test]
773	async fn a_share_wakes_when_a_sibling_goes_idle() {
774		let estimate = Producer::new();
775		let allocator = Allocator::new(estimate.consume());
776		estimate.set(Some(bps(2_000_000))).unwrap();
777
778		let (_mine_broadcast, mine) = track(VIDEO);
779		let _mine_sub = mine.consume();
780		let mine_share_reserved = allocator.reserve(&mine.demand(), bps(4_000_000));
781		let mut mine_share = mine_share_reserved.consumer();
782
783		let (_sibling_broadcast, sibling) = track(VIDEO);
784		let sibling_sub = sibling.consume();
785		let _sibling_share = allocator.reserve(&sibling.demand(), bps(4_000_000));
786
787		assert_eq!(mine_share.changed().await.unwrap(), Some(bps(1_000_000)));
788
789		// The sibling's last viewer leaves, so its half comes back to us.
790		drop(sibling_sub);
791		assert_eq!(mine_share.changed().await.unwrap(), Some(bps(2_000_000)));
792	}
793
794	/// Records whether a parked poll was actually woken, which re-reading state on a
795	/// fresh `poll` can't tell you: a lost wakeup still looks correct on the next poll
796	/// and only shows up as a task that never runs again.
797	#[derive(Default)]
798	struct Woken(std::sync::atomic::AtomicBool);
799
800	impl std::task::Wake for Woken {
801		fn wake(self: std::sync::Arc<Self>) {
802			self.wake_by_ref();
803		}
804
805		fn wake_by_ref(self: &std::sync::Arc<Self>) {
806			self.0.store(true, std::sync::atomic::Ordering::SeqCst);
807		}
808	}
809
810	impl Woken {
811		/// A flag and its waiter. Both are returned because a [`kio::WaiterList`]
812		/// holds only a `Weak`, so a waiter dropped at the end of the calling
813		/// statement takes its own registration with it and never fires.
814		fn new() -> (std::sync::Arc<Self>, kio::Waiter) {
815			let flag = std::sync::Arc::new(Self::default());
816			let waiter = kio::Waiter::new(std::task::Waker::from(flag.clone()));
817			(flag, waiter)
818		}
819
820		fn woken(&self) -> bool {
821			self.0.load(std::sync::atomic::Ordering::SeqCst)
822		}
823	}
824
825	/// Regression: a share whose slice didn't move still has to leave the estimate's
826	/// waker armed. Reservations cap the slice, so most estimate changes don't reach
827	/// it, and a poll that returned the value without arming anything would park the
828	/// encoder with nothing left to wake it.
829	#[tokio::test]
830	async fn an_unchanged_slice_keeps_watching_the_estimate() {
831		let estimate = Producer::new();
832		let allocator = Allocator::new(estimate.consume());
833
834		let (_broadcast, track) = track(VIDEO);
835		let _sub = track.consume();
836		let share_reserved = allocator.reserve(&track.demand(), bps(4_000_000));
837		let mut share = share_reserved.consumer();
838
839		estimate.set(Some(bps(10_000_000))).unwrap();
840		assert_eq!(share.changed().await.unwrap(), Some(bps(4_000_000)));
841
842		// Still miles above the reservation, so the slice holds at 4 Mbps: the share
843		// observes the change, decides it doesn't move, and parks.
844		let (woken, waiter) = Woken::new();
845		estimate.set(Some(bps(9_000_000))).unwrap();
846		assert!(share.poll_changed(&waiter).is_pending());
847
848		// The estimate finally drops past the reservation: this has to reach it.
849		estimate.set(Some(bps(1_000_000))).unwrap();
850		assert!(woken.woken(), "a parked share must be woken by the next estimate");
851		assert_eq!(share.changed().await.unwrap(), Some(bps(1_000_000)));
852	}
853
854	/// The same, for the other input: a sibling's demand.
855	#[tokio::test]
856	async fn a_parked_share_is_woken_by_sibling_demand() {
857		let estimate = Producer::new();
858		let allocator = Allocator::new(estimate.consume());
859		estimate.set(Some(bps(2_000_000))).unwrap();
860
861		let (_mine_broadcast, mine) = track(VIDEO);
862		let _mine_sub = mine.consume();
863		let share_reserved = allocator.reserve(&mine.demand(), bps(4_000_000));
864		let mut share = share_reserved.consumer();
865
866		let (_sibling_broadcast, sibling) = track(VIDEO);
867		let sibling_sub = sibling.consume();
868		let _sibling_share = allocator.reserve(&sibling.demand(), bps(4_000_000));
869
870		assert_eq!(share.changed().await.unwrap(), Some(bps(1_000_000)));
871
872		let (woken, waiter) = Woken::new();
873		assert!(share.poll_changed(&waiter).is_pending());
874		drop(sibling_sub);
875		assert!(woken.woken(), "a sibling going idle must wake a parked share");
876	}
877
878	/// A share follows the estimate's own lifecycle: unavailable while disconnected
879	/// (hold the current rate), terminal once the session is gone for good.
880	#[tokio::test]
881	async fn a_share_follows_the_estimate_lifecycle() {
882		let estimate = Producer::new();
883		let allocator = Allocator::new(estimate.consume());
884
885		let (_broadcast, track) = track(VIDEO);
886		let _sub = track.consume();
887		let share_reserved = allocator.reserve(&track.demand(), bps(4_000_000));
888		let mut share = share_reserved.consumer();
889
890		estimate.set(Some(bps(2_000_000))).unwrap();
891		assert_eq!(share.changed().await.unwrap(), Some(bps(2_000_000)));
892
893		estimate.set(None).unwrap();
894		assert_eq!(share.changed().await.unwrap(), None);
895
896		estimate.abort(Error::Cancel).unwrap();
897		assert!(share.changed().await.is_err());
898		assert!(share.changed().await.is_err());
899	}
900
901	/// A closed track's entry can't linger: it would be walked on every poll for the
902	/// rest of the connection, and a long-lived publisher churns tracks.
903	#[tokio::test]
904	async fn a_closed_track_is_pruned() {
905		let estimate = Producer::new();
906		let allocator = Allocator::new(estimate.consume());
907
908		// Both shares are held: dropping one releases its reservation on its own,
909		// which would prove nothing about pruning the track that closed.
910		let (_first_broadcast, first) = track(VIDEO);
911		let _first_share = allocator.reserve(&first.demand(), bps(4_000_000));
912		first.abort(Error::Cancel).unwrap();
913
914		let (_second_broadcast, second) = track(VIDEO);
915		let _second_share = allocator.reserve(&second.demand(), bps(4_000_000));
916
917		assert_eq!(allocator.registry.read().entries.len(), 1);
918	}
919
920	/// Dropping the reservation hands the room back. The registry only ever prunes
921	/// tracks that have *closed*, so a claim left behind by a track that is still
922	/// publishing would stand for the rest of the connection.
923	#[tokio::test]
924	async fn dropping_a_reservation_releases_it() {
925		let estimate = Producer::new();
926		let allocator = Allocator::new(estimate.consume());
927		estimate.set(Some(bps(2_000_000))).unwrap();
928
929		let (_first_broadcast, first) = track(VIDEO);
930		let _first_sub = first.consume();
931		let first_reserved = allocator.reserve(&first.demand(), bps(4_000_000));
932
933		let (_second_broadcast, second) = track(VIDEO);
934		let _second_sub = second.consume();
935		let second_reserved = allocator.reserve(&second.demand(), bps(4_000_000));
936		assert_eq!(second_reserved.peek(), Some(bps(1_000_000)));
937
938		// A read handle is not the claim: the reservation outliving it is what keeps the
939		// room, and the reservation going away is what returns it.
940		let mut orphan = first_reserved.consumer();
941		assert_eq!(orphan.changed().await.unwrap(), Some(bps(1_000_000)));
942
943		drop(first_reserved);
944		assert_eq!(allocator.registry.read().entries.len(), 1);
945		assert_eq!(second_reserved.peek(), Some(bps(2_000_000)));
946
947		// The orphaned reader is woken and told, rather than being left parked on a slice
948		// that will never move again. It reports the same "hold your rate" as an unwatched
949		// track, not a grant of zero that would tell an encoder to stop.
950		assert_eq!(orphan.changed().await.unwrap(), None);
951		assert_eq!(orphan.peek(), None);
952	}
953
954	/// A sender whose ceiling moved (a capture reopening at a resolution it negotiated
955	/// with the device) changes the claim in place. Re-reserving instead would claim
956	/// twice, since nothing releases the first entry while the track is still alive.
957	#[tokio::test]
958	async fn update_changes_the_claim_in_place() {
959		let estimate = Producer::new();
960		let allocator = Allocator::new(estimate.consume());
961		estimate.set(Some(bps(6_000_000))).unwrap();
962
963		let (_small_broadcast, small) = track(VIDEO);
964		let _small_sub = small.consume();
965		let small_reserved = allocator.reserve(&small.demand(), bps(1_000_000));
966
967		let (_large_broadcast, large) = track(VIDEO);
968		let _large_sub = large.consume();
969		let large_reserved = allocator.reserve(&large.demand(), bps(8_000_000));
970
971		// Max-min fair: the small claim is satisfied in full, the rest goes to the other.
972		assert_eq!(small_reserved.peek(), Some(bps(1_000_000)));
973		assert_eq!(large_reserved.peek(), Some(bps(5_000_000)));
974
975		// It reopens at a mode that can use much more, and the split follows without a
976		// second entry appearing.
977		small_reserved.update(bps(4_000_000));
978		assert_eq!(allocator.registry.read().entries.len(), 2);
979		assert_eq!(small_reserved.peek(), Some(bps(3_000_000)));
980		assert_eq!(large_reserved.peek(), Some(bps(3_000_000)));
981
982		// And back down, which has to release the difference rather than hold it.
983		small_reserved.update(bps(1_000_000));
984		assert_eq!(large_reserved.peek(), Some(bps(5_000_000)));
985	}
986
987	/// A reader parked on `changed` has to be woken by its own reservation moving, not
988	/// just by the estimate or a sibling: the encoder is sitting in a `select!` on it.
989	#[tokio::test]
990	async fn update_wakes_a_parked_reader() {
991		let estimate = Producer::new();
992		let allocator = Allocator::new(estimate.consume());
993		estimate.set(Some(bps(6_000_000))).unwrap();
994
995		let (_broadcast, producer) = track(VIDEO);
996		let _sub = producer.consume();
997		let reserved = allocator.reserve(&producer.demand(), bps(1_000_000));
998		let mut share = reserved.consumer();
999		assert_eq!(share.changed().await.unwrap(), Some(bps(1_000_000)));
1000
1001		let (woken, waiter) = Woken::new();
1002		assert!(share.poll_changed(&waiter).is_pending());
1003
1004		reserved.update(bps(4_000_000));
1005		assert!(woken.woken(), "raising the ceiling must wake the reader");
1006		assert_eq!(share.changed().await.unwrap(), Some(bps(4_000_000)));
1007	}
1008
1009	/// The registry is the allocator's, not a share's: a share that outlives every
1010	/// allocator reports the estimate as gone rather than keeping the channel open.
1011	#[tokio::test]
1012	async fn a_share_outliving_the_allocator_reports_closed() {
1013		let estimate = Producer::new();
1014		let allocator = Allocator::new(estimate.consume());
1015
1016		let (_broadcast, producer) = track(VIDEO);
1017		let _sub = producer.consume();
1018		let share_reserved = allocator.reserve(&producer.demand(), bps(4_000_000));
1019		let mut share = share_reserved.consumer();
1020
1021		drop(allocator);
1022		drop(estimate);
1023		assert!(share.changed().await.is_err());
1024	}
1025
1026	/// An unavailable estimate and a dead producer must not look alike: a caller
1027	/// holds its rate for the former and stops watching for the latter.
1028	/// Reporting closure as `Ok(None)` would spin any `select!` over `changed()`,
1029	/// because a closed channel is always immediately ready.
1030	#[tokio::test]
1031	async fn closed_is_distinct_from_unavailable() {
1032		let producer = Producer::new();
1033		let mut consumer = producer.consume();
1034
1035		producer.set(Some(bps(1_000_000))).unwrap();
1036		assert_eq!(consumer.changed().await.unwrap(), Some(bps(1_000_000)));
1037
1038		// Live, but the estimate went away (e.g. disconnected): still watchable.
1039		producer.set(None).unwrap();
1040		assert_eq!(consumer.changed().await.unwrap(), None);
1041
1042		// Gone for good.
1043		producer.abort(Error::Cancel).unwrap();
1044		assert!(consumer.changed().await.is_err());
1045		// And it stays terminal rather than flapping back to a value.
1046		assert!(consumer.changed().await.is_err());
1047	}
1048}