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 carries a publisher that leaves `priority` at its default:
236	/// one tier of audio and video still serves audio's small reservation in full
237	/// 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	/// A publisher that doesn't set [`track::Info::priority`] puts audio and video
667	/// in one tier. That has to come out right anyway, and it does: max-min fair
668	/// satisfies the small claim first, so audio still gets its full reservation
669	/// and video takes the rest. Priority only changes the answer once a tier's
670	/// smaller claims outgrow an even split.
671	#[test]
672	fn one_tier_still_serves_audio_before_video() {
673		let flat = [want(0, 0, 128_000), want(1, 0, 4_000_000)];
674		let tiered = [want(0, AUDIO, 128_000), want(1, VIDEO, 4_000_000)];
675
676		for wants in [flat, tiered] {
677			assert_eq!(allocate(bps(2_000_000), &wants, 0), Some(bps(128_000)));
678			assert_eq!(allocate(bps(2_000_000), &wants, 1), Some(bps(1_872_000)));
679		}
680	}
681
682	#[test]
683	fn an_even_tier_splits_evenly() {
684		let wants = [want(0, VIDEO, 4_000_000), want(1, VIDEO, 4_000_000)];
685
686		assert_eq!(allocate(bps(6_000_000), &wants, 0), Some(bps(3_000_000)));
687		assert_eq!(allocate(bps(6_000_000), &wants, 1), Some(bps(3_000_000)));
688	}
689
690	/// Max-min fair, not an even split: a 360p rung sharing with a 1080p one takes
691	/// only what it asked for and leaves the rest, instead of both being held to half.
692	#[test]
693	fn a_small_share_frees_what_it_does_not_want() {
694		let wants = [want(0, VIDEO, 1_000_000), want(1, VIDEO, 8_000_000)];
695
696		assert_eq!(allocate(bps(6_000_000), &wants, 0), Some(bps(1_000_000)));
697		assert_eq!(allocate(bps(6_000_000), &wants, 1), Some(bps(5_000_000)));
698	}
699
700	/// Capping at the reservation is what keeps an uncongested link encoding at
701	/// exactly the configured rate. Spreading the surplus instead would only matter
702	/// if senders scaled a fraction of their grant, which is what headroom did.
703	#[test]
704	fn surplus_is_left_unclaimed() {
705		assert_eq!(
706			allocate(bps(10_000_000), &[want(0, VIDEO, 4_000_000)], 0),
707			Some(bps(4_000_000))
708		);
709	}
710
711	#[test]
712	fn an_unregistered_share_has_no_grant() {
713		assert_eq!(allocate(bps(1_000_000), &[], 0), None);
714		assert_eq!(allocate(bps(1_000_000), &[want(0, VIDEO, 1_000)], 7), None);
715	}
716
717	/// The headline case: two encoders on one connection must not each target the
718	/// whole estimate.
719	#[tokio::test]
720	async fn concurrent_tracks_split_the_estimate() {
721		let estimate = Producer::new();
722		let allocator = Allocator::new(estimate.consume());
723
724		let (_first_broadcast, first) = track(VIDEO);
725		let _first_sub = first.consume();
726		let first = allocator.reserve(&first.demand(), bps(4_000_000));
727		let mut first_share = first.consumer();
728
729		estimate.set(Some(bps(2_000_000))).unwrap();
730		// Alone, it gets everything it asked for that the link can carry.
731		assert_eq!(first_share.changed().await.unwrap(), Some(bps(2_000_000)));
732
733		// A second encoder starts. The first must give half back without anyone
734		// telling it to, or the two together target 200% of the uplink.
735		let (_second_broadcast, second) = track(VIDEO);
736		let _second_sub = second.consume();
737		let second_share = allocator.reserve(&second.demand(), bps(4_000_000));
738
739		assert_eq!(first_share.changed().await.unwrap(), Some(bps(1_000_000)));
740		assert_eq!(second_share.peek(), Some(bps(1_000_000)));
741	}
742
743	/// An unwatched track releases its reservation, since `publish_capture` stops
744	/// encoding entirely while nothing is subscribed.
745	#[tokio::test]
746	async fn an_idle_track_claims_nothing() {
747		let estimate = Producer::new();
748		let allocator = Allocator::new(estimate.consume());
749		estimate.set(Some(bps(2_000_000))).unwrap();
750
751		let (_watched_broadcast, watched) = track(VIDEO);
752		let _watched_sub = watched.consume();
753		let watched_share = allocator.reserve(&watched.demand(), bps(4_000_000));
754
755		let (_idle_broadcast, idle) = track(VIDEO);
756		let idle_share = allocator.reserve(&idle.demand(), bps(4_000_000));
757
758		assert_eq!(watched_share.peek(), Some(bps(2_000_000)));
759		// Not zero: an idle share reports "no opinion" so a sender that is mid-shutdown
760		// holds its rate instead of retuning to the floor on the way out.
761		assert_eq!(idle_share.peek(), None);
762
763		// It joins, and the split happens.
764		let _idle_sub = idle.consume();
765		assert_eq!(watched_share.peek(), Some(bps(1_000_000)));
766		assert_eq!(idle_share.peek(), Some(bps(1_000_000)));
767	}
768
769	/// Demand transitions have to wake a parked share, not just change what a later
770	/// `peek` would see; the encoder is sitting in a `select!` on `changed`.
771	#[tokio::test]
772	async fn a_share_wakes_when_a_sibling_goes_idle() {
773		let estimate = Producer::new();
774		let allocator = Allocator::new(estimate.consume());
775		estimate.set(Some(bps(2_000_000))).unwrap();
776
777		let (_mine_broadcast, mine) = track(VIDEO);
778		let _mine_sub = mine.consume();
779		let mine_share_reserved = allocator.reserve(&mine.demand(), bps(4_000_000));
780		let mut mine_share = mine_share_reserved.consumer();
781
782		let (_sibling_broadcast, sibling) = track(VIDEO);
783		let sibling_sub = sibling.consume();
784		let _sibling_share = allocator.reserve(&sibling.demand(), bps(4_000_000));
785
786		assert_eq!(mine_share.changed().await.unwrap(), Some(bps(1_000_000)));
787
788		// The sibling's last viewer leaves, so its half comes back to us.
789		drop(sibling_sub);
790		assert_eq!(mine_share.changed().await.unwrap(), Some(bps(2_000_000)));
791	}
792
793	/// Records whether a parked poll was actually woken, which re-reading state on a
794	/// fresh `poll` can't tell you: a lost wakeup still looks correct on the next poll
795	/// and only shows up as a task that never runs again.
796	#[derive(Default)]
797	struct Woken(std::sync::atomic::AtomicBool);
798
799	impl std::task::Wake for Woken {
800		fn wake(self: std::sync::Arc<Self>) {
801			self.wake_by_ref();
802		}
803
804		fn wake_by_ref(self: &std::sync::Arc<Self>) {
805			self.0.store(true, std::sync::atomic::Ordering::SeqCst);
806		}
807	}
808
809	impl Woken {
810		/// A flag and its waiter. Both are returned because a [`kio::WaiterList`]
811		/// holds only a `Weak`, so a waiter dropped at the end of the calling
812		/// statement takes its own registration with it and never fires.
813		fn new() -> (std::sync::Arc<Self>, kio::Waiter) {
814			let flag = std::sync::Arc::new(Self::default());
815			let waiter = kio::Waiter::new(std::task::Waker::from(flag.clone()));
816			(flag, waiter)
817		}
818
819		fn woken(&self) -> bool {
820			self.0.load(std::sync::atomic::Ordering::SeqCst)
821		}
822	}
823
824	/// Regression: a share whose slice didn't move still has to leave the estimate's
825	/// waker armed. Reservations cap the slice, so most estimate changes don't reach
826	/// it, and a poll that returned the value without arming anything would park the
827	/// encoder with nothing left to wake it.
828	#[tokio::test]
829	async fn an_unchanged_slice_keeps_watching_the_estimate() {
830		let estimate = Producer::new();
831		let allocator = Allocator::new(estimate.consume());
832
833		let (_broadcast, track) = track(VIDEO);
834		let _sub = track.consume();
835		let share_reserved = allocator.reserve(&track.demand(), bps(4_000_000));
836		let mut share = share_reserved.consumer();
837
838		estimate.set(Some(bps(10_000_000))).unwrap();
839		assert_eq!(share.changed().await.unwrap(), Some(bps(4_000_000)));
840
841		// Still miles above the reservation, so the slice holds at 4 Mbps: the share
842		// observes the change, decides it doesn't move, and parks.
843		let (woken, waiter) = Woken::new();
844		estimate.set(Some(bps(9_000_000))).unwrap();
845		assert!(share.poll_changed(&waiter).is_pending());
846
847		// The estimate finally drops past the reservation: this has to reach it.
848		estimate.set(Some(bps(1_000_000))).unwrap();
849		assert!(woken.woken(), "a parked share must be woken by the next estimate");
850		assert_eq!(share.changed().await.unwrap(), Some(bps(1_000_000)));
851	}
852
853	/// The same, for the other input: a sibling's demand.
854	#[tokio::test]
855	async fn a_parked_share_is_woken_by_sibling_demand() {
856		let estimate = Producer::new();
857		let allocator = Allocator::new(estimate.consume());
858		estimate.set(Some(bps(2_000_000))).unwrap();
859
860		let (_mine_broadcast, mine) = track(VIDEO);
861		let _mine_sub = mine.consume();
862		let share_reserved = allocator.reserve(&mine.demand(), bps(4_000_000));
863		let mut share = share_reserved.consumer();
864
865		let (_sibling_broadcast, sibling) = track(VIDEO);
866		let sibling_sub = sibling.consume();
867		let _sibling_share = allocator.reserve(&sibling.demand(), bps(4_000_000));
868
869		assert_eq!(share.changed().await.unwrap(), Some(bps(1_000_000)));
870
871		let (woken, waiter) = Woken::new();
872		assert!(share.poll_changed(&waiter).is_pending());
873		drop(sibling_sub);
874		assert!(woken.woken(), "a sibling going idle must wake a parked share");
875	}
876
877	/// A share follows the estimate's own lifecycle: unavailable while disconnected
878	/// (hold the current rate), terminal once the session is gone for good.
879	#[tokio::test]
880	async fn a_share_follows_the_estimate_lifecycle() {
881		let estimate = Producer::new();
882		let allocator = Allocator::new(estimate.consume());
883
884		let (_broadcast, track) = track(VIDEO);
885		let _sub = track.consume();
886		let share_reserved = allocator.reserve(&track.demand(), bps(4_000_000));
887		let mut share = share_reserved.consumer();
888
889		estimate.set(Some(bps(2_000_000))).unwrap();
890		assert_eq!(share.changed().await.unwrap(), Some(bps(2_000_000)));
891
892		estimate.set(None).unwrap();
893		assert_eq!(share.changed().await.unwrap(), None);
894
895		estimate.abort(Error::Cancel).unwrap();
896		assert!(share.changed().await.is_err());
897		assert!(share.changed().await.is_err());
898	}
899
900	/// A closed track's entry can't linger: it would be walked on every poll for the
901	/// rest of the connection, and a long-lived publisher churns tracks.
902	#[tokio::test]
903	async fn a_closed_track_is_pruned() {
904		let estimate = Producer::new();
905		let allocator = Allocator::new(estimate.consume());
906
907		// Both shares are held: dropping one releases its reservation on its own,
908		// which would prove nothing about pruning the track that closed.
909		let (_first_broadcast, first) = track(VIDEO);
910		let _first_share = allocator.reserve(&first.demand(), bps(4_000_000));
911		first.abort(Error::Cancel).unwrap();
912
913		let (_second_broadcast, second) = track(VIDEO);
914		let _second_share = allocator.reserve(&second.demand(), bps(4_000_000));
915
916		assert_eq!(allocator.registry.read().entries.len(), 1);
917	}
918
919	/// Dropping the reservation hands the room back. The registry only ever prunes
920	/// tracks that have *closed*, so a claim left behind by a track that is still
921	/// publishing would stand for the rest of the connection.
922	#[tokio::test]
923	async fn dropping_a_reservation_releases_it() {
924		let estimate = Producer::new();
925		let allocator = Allocator::new(estimate.consume());
926		estimate.set(Some(bps(2_000_000))).unwrap();
927
928		let (_first_broadcast, first) = track(VIDEO);
929		let _first_sub = first.consume();
930		let first_reserved = allocator.reserve(&first.demand(), bps(4_000_000));
931
932		let (_second_broadcast, second) = track(VIDEO);
933		let _second_sub = second.consume();
934		let second_reserved = allocator.reserve(&second.demand(), bps(4_000_000));
935		assert_eq!(second_reserved.peek(), Some(bps(1_000_000)));
936
937		// A read handle is not the claim: the reservation outliving it is what keeps the
938		// room, and the reservation going away is what returns it.
939		let mut orphan = first_reserved.consumer();
940		assert_eq!(orphan.changed().await.unwrap(), Some(bps(1_000_000)));
941
942		drop(first_reserved);
943		assert_eq!(allocator.registry.read().entries.len(), 1);
944		assert_eq!(second_reserved.peek(), Some(bps(2_000_000)));
945
946		// The orphaned reader is woken and told, rather than being left parked on a slice
947		// that will never move again. It reports the same "hold your rate" as an unwatched
948		// track, not a grant of zero that would tell an encoder to stop.
949		assert_eq!(orphan.changed().await.unwrap(), None);
950		assert_eq!(orphan.peek(), None);
951	}
952
953	/// A sender whose ceiling moved (a capture reopening at a resolution it negotiated
954	/// with the device) changes the claim in place. Re-reserving instead would claim
955	/// twice, since nothing releases the first entry while the track is still alive.
956	#[tokio::test]
957	async fn update_changes_the_claim_in_place() {
958		let estimate = Producer::new();
959		let allocator = Allocator::new(estimate.consume());
960		estimate.set(Some(bps(6_000_000))).unwrap();
961
962		let (_small_broadcast, small) = track(VIDEO);
963		let _small_sub = small.consume();
964		let small_reserved = allocator.reserve(&small.demand(), bps(1_000_000));
965
966		let (_large_broadcast, large) = track(VIDEO);
967		let _large_sub = large.consume();
968		let large_reserved = allocator.reserve(&large.demand(), bps(8_000_000));
969
970		// Max-min fair: the small claim is satisfied in full, the rest goes to the other.
971		assert_eq!(small_reserved.peek(), Some(bps(1_000_000)));
972		assert_eq!(large_reserved.peek(), Some(bps(5_000_000)));
973
974		// It reopens at a mode that can use much more, and the split follows without a
975		// second entry appearing.
976		small_reserved.update(bps(4_000_000));
977		assert_eq!(allocator.registry.read().entries.len(), 2);
978		assert_eq!(small_reserved.peek(), Some(bps(3_000_000)));
979		assert_eq!(large_reserved.peek(), Some(bps(3_000_000)));
980
981		// And back down, which has to release the difference rather than hold it.
982		small_reserved.update(bps(1_000_000));
983		assert_eq!(large_reserved.peek(), Some(bps(5_000_000)));
984	}
985
986	/// A reader parked on `changed` has to be woken by its own reservation moving, not
987	/// just by the estimate or a sibling: the encoder is sitting in a `select!` on it.
988	#[tokio::test]
989	async fn update_wakes_a_parked_reader() {
990		let estimate = Producer::new();
991		let allocator = Allocator::new(estimate.consume());
992		estimate.set(Some(bps(6_000_000))).unwrap();
993
994		let (_broadcast, producer) = track(VIDEO);
995		let _sub = producer.consume();
996		let reserved = allocator.reserve(&producer.demand(), bps(1_000_000));
997		let mut share = reserved.consumer();
998		assert_eq!(share.changed().await.unwrap(), Some(bps(1_000_000)));
999
1000		let (woken, waiter) = Woken::new();
1001		assert!(share.poll_changed(&waiter).is_pending());
1002
1003		reserved.update(bps(4_000_000));
1004		assert!(woken.woken(), "raising the ceiling must wake the reader");
1005		assert_eq!(share.changed().await.unwrap(), Some(bps(4_000_000)));
1006	}
1007
1008	/// The registry is the allocator's, not a share's: a share that outlives every
1009	/// allocator reports the estimate as gone rather than keeping the channel open.
1010	#[tokio::test]
1011	async fn a_share_outliving_the_allocator_reports_closed() {
1012		let estimate = Producer::new();
1013		let allocator = Allocator::new(estimate.consume());
1014
1015		let (_broadcast, producer) = track(VIDEO);
1016		let _sub = producer.consume();
1017		let share_reserved = allocator.reserve(&producer.demand(), bps(4_000_000));
1018		let mut share = share_reserved.consumer();
1019
1020		drop(allocator);
1021		drop(estimate);
1022		assert!(share.changed().await.is_err());
1023	}
1024
1025	/// An unavailable estimate and a dead producer must not look alike: a caller
1026	/// holds its rate for the former and stops watching for the latter.
1027	/// Reporting closure as `Ok(None)` would spin any `select!` over `changed()`,
1028	/// because a closed channel is always immediately ready.
1029	#[tokio::test]
1030	async fn closed_is_distinct_from_unavailable() {
1031		let producer = Producer::new();
1032		let mut consumer = producer.consume();
1033
1034		producer.set(Some(bps(1_000_000))).unwrap();
1035		assert_eq!(consumer.changed().await.unwrap(), Some(bps(1_000_000)));
1036
1037		// Live, but the estimate went away (e.g. disconnected): still watchable.
1038		producer.set(None).unwrap();
1039		assert_eq!(consumer.changed().await.unwrap(), None);
1040
1041		// Gone for good.
1042		producer.abort(Error::Cancel).unwrap();
1043		assert!(consumer.changed().await.is_err());
1044		// And it stays terminal rather than flapping back to a value.
1045		assert!(consumer.changed().await.is_err());
1046	}
1047}