fix(mux): bound what the decode clock holds in MKV and TS export

rs/moq-mux/Cargo.toml

@@ -65,6 +65,10 @@ harness = false
 name = "jitter"
 harness = false
 
+[[bench]]
+name = "decode_clock"
+harness = false
+
 [[bench]]
 name = "containers"
 harness = false

rs/moq-mux/benches/decode_clock.rs

@@ -0,0 +1,98 @@
+//! The decode clock swept over how many frames it holds and how their PTS run: draining what
+//! it holds, and settling frames as they are read until the overdue rule bounds them.
+use std::time::Duration;
+
+use bytes::Bytes;
+use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
+use moq_clock::Instant;
+
+#[path = "../src/container/decode.rs"]
+#[allow(dead_code)]
+mod decode;
+
+/// A 40 ms frame in 90 kHz ticks.
+const FRAME: u64 = 3_600;
+
+/// The `index`-th frame's PTS in decode order: stuck, running backward, or a valid B-pyramid.
+fn pts(shape: &str, index: u64) -> u64 {
+	let base = 1_000 * 90_000;
+	match shape {
+		"flat" => base,
+		"backward" => base - index,
+		_ => {
+			const PYRAMID: [u64; 9] = [0, 4, 2, 1, 3, 8, 6, 5, 7];
+			base + (index / 9 * 9 + PYRAMID[(index % 9) as usize]) * FRAME
+		}
+	}
+}
+
+/// Flush `held` frames read with nothing settled, as at a track's end.
+fn drain(c: &mut Criterion) {
+	let mut group = c.benchmark_group("decode_clock_drain");
+	group.measurement_time(Duration::from_secs(1));
+	group.warm_up_time(Duration::from_millis(200));
+	group.sample_size(10);
+	for shape in ["flat", "backward", "pyramid"] {
+		for held in [1_000u64, 10_000, 100_000] {
+			for size in [64usize, 4_096] {
+				let payload = Bytes::from(vec![0; size]);
+				group.throughput(Throughput::Elements(held));
+				group.bench_function(BenchmarkId::new(format!("{shape}/{size}B"), held), |b| {
+					b.iter(|| {
+						let mut clock = decode::DecodeClock::default();
+						for index in 0..held {
+							clock.push(payload.clone(), pts(shape, index), Instant::ZERO);
+						}
+						let mut drained = 0;
+						while clock.pop(decode::MAX_REORDER, 0, decode::Settle::All).is_some() {
+							drained += 1;
+						}
+						assert_eq!(drained, held);
+					})
+				});
+			}
+		}
+	}
+	group.finish();
+}
+
+/// Read `frames` frames 1 ms apart, a 120 ms lookahead and 500 ms delay, settling each as an
+/// exporter does: flat or backward PTS holds what the overdue rule lets through.
+fn read(c: &mut Criterion) {
+	let mut group = c.benchmark_group("decode_clock_read");
+	group.measurement_time(Duration::from_secs(1));
+	group.warm_up_time(Duration::from_millis(200));
+	group.sample_size(10);
+	let lookahead = 120 * 90;
+	let delay = Duration::from_millis(500);
+	for shape in ["flat", "backward", "pyramid"] {
+		for frames in [1_000u64, 10_000, 100_000] {
+			group.throughput(Throughput::Elements(frames));
+			group.bench_function(BenchmarkId::new(shape, frames), |b| {
+				b.iter(|| {
+					let mut clock = decode::DecodeClock::default();
+					let (mut held, mut most) = (0, 0);
+					for index in 0..frames {
+						let latest = Instant::ZERO + Duration::from_millis(index);
+						clock.push(index, pts(shape, index), latest);
+						held += 1;
+						loop {
+							let settle = clock.settle(lookahead, latest, delay);
+							if clock.pop(lookahead, 0, settle).is_none() {
+								break;
+							}
+							held -= 1;
+						}
+						most = most.max(held);
+					}
+					// What arrives in the deepest lookahead and the delay, a frame a millisecond.
+					assert!(most <= 2_501, "{shape}: {most} held");
+				})
+			});
+		}
+	}
+	group.finish();
+}
+
+criterion_group!(benches, drain, read);
+criterion_main!(benches);

rs/moq-mux/src/container/decode.rs

@@ -1,8 +1,22 @@
 use std::collections::VecDeque;
+use std::time::Duration;
+
+use moq_clock::Instant;
 
 /// The largest supported reorder span in 90 kHz ticks.
 pub(crate) const MAX_REORDER: u64 = 2 * 90_000;
 
+/// Which frames [`DecodeClock::pop`] gives out.
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub(crate) enum Settle {
+	/// Those whose slot the highest PTS read settled.
+	Due,
+	/// Those, and the frame held longest, which is overdue ([`DecodeClock::settle`]).
+	Front,
+	/// Every frame held, as nothing still to come is presented before them.
+	All,
+}
+
 /// A video rendition's decode clock, authored from presentation timestamps.
 ///
 /// [`crate::container::Frame`] carries only a presentation timestamp and frames reach the muxer in decode order,
@@ -23,13 +37,21 @@ pub(crate) const MAX_REORDER: u64 = 2 * 90_000;
 /// (reordering deeper than all of those) is nudged one tick past the last DTS, so the clock
 /// stays strictly increasing.
 ///
+/// A source whose PTS stops advancing would hold every frame after it here, so a frame that
+/// has waited longer than its lookahead plus the export's delay is overdue, and settled at
+/// once ([`Self::settle`]): a live source's PTS keeps pace with its arrival, so the frames
+/// held stay within what the source sends in that time, like the jitter buffer's.
+///
 /// Ticks are continuous (unwrapped) 90 kHz, so the clock never wraps mid-stream; the 33-bit
 /// wire wrap happens once at emission.
 pub(crate) struct DecodeClock<T> {
-	/// Frames read but not yet given a DTS, in decode order, with their PTS.
-	held: VecDeque<(T, u64)>,
+	/// Frames read but not yet given a DTS, in decode order, with their PTS and arrival.
+	held: VecDeque<(T, u64, Instant)>,
 	/// The held frames' PTS, ascending: the slots they decode in, in turn.
-	slots: Vec<u64>,
+	slots: VecDeque<u64>,
+	/// How many leading slots, paired with the held frame of the same rank, already bound the
+	/// delay, so draining a backlog scans each pair once.
+	bounded: usize,
 	/// The highest PTS read.
 	high: Option<u64>,
 	/// The furthest a frame was presented below the highest PTS read before it, plus a tick.
@@ -44,7 +66,8 @@ impl<T> Default for DecodeClock<T> {
 	fn default() -> Self {
 		Self {
 			held: VecDeque::new(),
-			slots: Vec::new(),
+			slots: VecDeque::new(),
+			bounded: 0,
 			high: None,
 			reach: 0,
 			delay: 0,
@@ -54,8 +77,8 @@ impl<T> Default for DecodeClock<T> {
 }
 
 impl<T> DecodeClock<T> {
-	/// Hold the next frame in decode order, presented at `pts`.
-	pub(crate) fn push(&mut self, item: T, pts: u64) {
+	/// Hold the next frame in decode order, presented at `pts`, which `arrived` when read.
+	pub(crate) fn push(&mut self, item: T, pts: u64, arrived: Instant) {
 		if let Some(high) = self.high
 			&& pts < high
 		{
@@ -64,30 +87,39 @@ impl<T> DecodeClock<T> {
 		self.high = Some(self.high.map_or(pts, |high| high.max(pts)));
 		let at = self.slots.partition_point(|&slot| slot <= pts);
 		self.slots.insert(at, pts);
-		self.held.push_back((item, pts));
+		self.bounded = self.bounded.min(at);
+		self.held.push_back((item, pts, arrived));
 	}
 
 	/// The next frame and its DTS, once its slot is settled `lookahead` past (the catalog
-	/// `jitter`), or at once when `flush`. `floor` is the least reorder delay, in ticks.
-	pub(crate) fn pop(&mut self, lookahead: u64, floor: u64, flush: bool) -> Option<(T, u64)> {
+	/// `jitter`), or at once as `settle` says. `floor` is the least reorder delay, in ticks.
+	pub(crate) fn pop(&mut self, lookahead: u64, floor: u64, settle: Settle) -> Option<(T, u64)> {
 		let high = self.high?;
-		let lookahead = lookahead.max(self.reach);
-		let settled = |slot: u64| flush || high >= slot.saturating_add(lookahead);
-		if !settled(*self.slots.first()?) {
+		let lookahead = self.lookahead(lookahead);
+		let front = *self.slots.front()?;
+		let settled = |slot: u64| match settle {
+			Settle::All => true,
+			Settle::Front if slot <= front => true,
+			_ => high >= slot.saturating_add(lookahead),
+		};
+		if !settled(front) {
 			return None;
 		}
 		// Every settled frame bounds the delay, so one decoding later never steps it back.
 		let mut delay = self.delay.max(floor);
-		for (&slot, (_, pts)) in self.slots.iter().zip(&self.held) {
+		let mut bounded = self.bounded;
+		for (&slot, (_, pts, _)) in self.slots.range(bounded..).zip(self.held.range(bounded..)) {
 			if !settled(slot) {
 				break;
 			}
 			delay = delay.max(slot.saturating_sub(*pts));
+			bounded += 1;
 		}
 		self.delay = delay.min(MAX_REORDER);
+		self.bounded = bounded - 1;
 
-		let slot = self.slots.remove(0);
-		let (item, _) = self.held.pop_front()?;
+		let slot = self.slots.pop_front()?;
+		let (item, ..) = self.held.pop_front()?;
 		let dts = slot.saturating_sub(self.delay);
 		let dts = match self.last {
 			Some(last) if dts <= last => last + 1,
@@ -96,4 +128,29 @@ impl<T> DecodeClock<T> {
 		self.last = Some(dts);
 		Some((item, dts))
 	}
+
+	/// What [`Self::pop`] settles short of a flush: [`Settle::Front`] once the frame held
+	/// longest is overdue, having waited, by the `latest` arrival read, longer than `lookahead`
+	/// plus the export's `delay`, else [`Settle::Due`].
+	pub(crate) fn settle(&self, lookahead: u64, latest: Instant, delay: Duration) -> Settle {
+		let Some((.., arrived)) = self.held.front() else {
+			return Settle::Due;
+		};
+		let lookahead = Duration::from_nanos(self.lookahead(lookahead) * 1_000_000_000 / 90_000);
+		match latest.saturating_duration_since(*arrived) > lookahead + delay {
+			true => Settle::Front,
+			false => Settle::Due,
+		}
+	}
+
+	/// How many frames are held.
+	#[cfg(test)]
+	pub(crate) fn len(&self) -> usize {
+		self.held.len()
+	}
+
+	/// How far past a slot the highest PTS must be to settle it, under `lookahead`.
+	fn lookahead(&self, lookahead: u64) -> u64 {
+		lookahead.max(self.reach).min(MAX_REORDER)
+	}
 }

rs/moq-mux/src/container/mkv/export.rs

@@ -13,6 +13,7 @@ use crate::catalog::Stream;
 use crate::container::ExportSource;
 use crate::container::Frame;
 use crate::container::consumer::Event;
+use crate::container::decode::{DecodeClock, Settle};
 use crate::container::mkv::Error;
 
 /// Matroska TimestampScale: 1 ms (in nanoseconds).
@@ -72,7 +73,7 @@ pub struct Export<S: Stream> {
 }
 
 struct MkvTrack {
-	clock: crate::container::decode::DecodeClock<crate::jitter::Arrival<Frame>>,
+	clock: DecodeClock<crate::jitter::Arrival<Frame>>,
 	config: Option<VideoConfig>,
 	restart: u64,
 	reorder: Option<(u64, u64)>,
@@ -176,7 +177,14 @@ impl MkvTrack {
 				.min(limit);
 			(floor.min(limit), lookahead)
 		});
-		while let Some((mut arrival, dts)) = self.clock.pop(lookahead, floor, flush) {
+		loop {
+			let settle = match flush {
+				true => Settle::All,
+				false => self.clock.settle(lookahead, self.latest, jitter.delay()),
+			};
+			let Some((mut arrival, dts)) = self.clock.pop(lookahead, floor, settle) else {
+				break;
+			};
 			if !self.finished && !flush {
 				arrival.arrived = arrival.arrived.max(self.latest);
 			}
@@ -404,7 +412,7 @@ impl<S: Stream> Export<S> {
 	/// How many frames the export holds between its sources and the output.
 	#[cfg(test)]
 	pub(super) fn buffered(&self) -> usize {
-		self.jitter.len()
+		self.jitter.len() + self.tracks.values().map(|track| track.clock.len()).sum::<usize>()
 	}
 
 	#[cfg(test)]
@@ -482,7 +490,8 @@ impl<S: Stream> Export<S> {
 						track.restart = arrival.restart;
 					}
 					let pts = (arrival.item.timestamp.as_nanos() * 90_000 / 1_000_000_000) as u64;
-					track.clock.push(arrival, pts);
+					let arrived = arrival.arrived;
+					track.clock.push(arrival, pts, arrived);
 					track.release(name, &mut self.jitter, false)?;
 				} else {
 					self.jitter.push(name.clone(), arrival)?;

rs/moq-mux/src/container/mkv/export_test.rs

@@ -1241,3 +1241,82 @@ async fn disabled_renditions_stay_out_of_the_header_and_idle_after_it() {
 	live.track.write(raw_frame(20_000, b"main-1", true)).unwrap();
 	assert!(chunk_now(&mut export).await.windows(6).any(|part| part == b"main-1"));
 }
+
+/// Video whose PTS stops advancing past the reorder lookahead in an open group: the export
+/// holds no more than what arrives within the lookahead and the delay. Flat PTS waits the
+/// catalog `jitter`; PTS running backward reads as reordering up to the deepest supported.
+#[tokio::test(start_paused = true)]
+async fn stalled_presentation_time_keeps_the_export_bounded() {
+	use std::time::Duration;
+	let step = Duration::from_millis(20);
+	let delay = Duration::from_millis(100);
+	let jitter = Duration::from_millis(120);
+	let deepest = Duration::from_secs(2);
+	let flat: fn(u64) -> u64 = |_| 10_000_000;
+	let backward: fn(u64) -> u64 = |index| 10_000_000 - index * 10_000;
+	for (name, pts, lookahead) in [("flat", flat, jitter), ("backward", backward, deepest)] {
+		let mut live = Live::avc3();
+		for config in live.catalog.modify().unwrap().video.renditions.values_mut() {
+			config.jitter = Some(jitter);
+		}
+		let mut export = super::Export::new(live.source().await, live.catalog_stream().await)
+			.with_delay(delay)
+			.with_fragment_duration(Duration::ZERO);
+		live.track.write(video_frame(pts(0), true)).unwrap();
+		export.next().await.unwrap().unwrap();
+		let mut most = 0;
+		for index in 1..500 {
+			tokio::time::advance(step).await;
+			live.track.write(video_frame(pts(index), false)).unwrap();
+			while let std::task::Poll::Ready(chunk) = export.poll_next(&kio::Waiter::noop()) {
+				chunk.unwrap().unwrap();
+			}
+			most = most.max(export.buffered());
+		}
+		// What arrived within the wait, the frame just read, and one due to go out.
+		let bound = ((lookahead + delay).as_millis() / step.as_millis()) as usize + 2;
+		assert!(most <= bound, "{name}: {most} frames buffered, over {bound}");
+	}
+}
+
+/// Live B-frames, read as they are sent, keep every frame and its presentation order: the
+/// bound on what the decode clock holds never cuts a valid reorder short.
+#[tokio::test(start_paused = true)]
+async fn live_reordered_video_keeps_every_frame() {
+	use std::time::Duration;
+	let frame = 40_000;
+	let mut live = Live::avc3();
+	for config in live.catalog.modify().unwrap().video.renditions.values_mut() {
+		config.jitter = Some(Duration::from_millis(120));
+	}
+	let mut export = super::Export::new(live.source().await, live.catalog_stream().await)
+		.with_delay(Duration::from_millis(500))
+		.with_fragment_duration(Duration::ZERO);
+	// A B-pyramid in decode order, a group of nine 40 ms frames at a time.
+	let pyramid = [0, 4, 2, 1, 3, 8, 6, 5, 7];
+	let order: Vec<u64> = (0..10u64)
+		.flat_map(|group| pyramid.map(|i| 1_000_000 + (group * 9 + i) * frame))
+		.collect();
+	live.track.write(video_frame(order[0], true)).unwrap();
+	export.next().await.unwrap().unwrap();
+	let mut blocks = 0;
+	let mut count = |chunk: Bytes| {
+		let mut cursor = Cursor::new(chunk);
+		blocks += WebmIterator::new(&mut cursor, &[])
+			.filter(|tag| matches!(tag, Ok(MatroskaSpec::SimpleBlock(_))))
+			.count();
+	};
+	for (index, &pts) in order.iter().enumerate().skip(1) {
+		tokio::time::advance(Duration::from_micros(frame)).await;
+		live.track.write(video_frame(pts, index % 9 == 0)).unwrap();
+		while let std::task::Poll::Ready(chunk) = export.poll_next(&kio::Waiter::noop()) {
+			count(chunk.unwrap().unwrap());
+		}
+	}
+	live.track.finish().unwrap();
+	live.catalog.finish().unwrap();
+	while let Some(chunk) = export.next().await.unwrap() {
+		count(chunk);
+	}
+	assert_eq!(blocks, order.len());
+}