mux: suppress adjacent identical PSI runs

.github/workflows/nightly.yml

@@ -169,6 +169,10 @@ jobs:
         if: ${{ !cancelled() }}
         run: nix develop --command cargo bench --locked -p moq-mux --bench catalog -- --test
 
+      - name: TS layout benchmark smoke
+        if: ${{ !cancelled() }}
+        run: nix develop --command just rs bench-ts-layout --test
+
       - name: RTMP chunk-stream benchmark smoke
         if: ${{ !cancelled() }}
         run: nix develop --command just rs bench-rtmp --test

quest/m1/tstd/psi-runs.md

@@ -24,3 +24,9 @@ own tables.
   after the fix. The AAC tune-in test still passes.
 
 Public API: none. Wire: none.
+
+### Remaining system-buffer guarantee
+
+The adjacent-copy implementation passes its original 3/4/8-video regression at 25/50/100 Mb/s, including AAC tune-in and changed-version ordering. It does not establish the broader system-buffer guarantee. Equal-weight aligned units still produce distinct tune-in runs separated by media: four video PIDs carrying 100 packets each reach 518.88 bytes on PAT at 25 Mb/s, and eight PIDs carrying 10 packets each reach 1022.72 bytes. Both fit their slot budget and exceed the 512-byte transport buffer after adjacent-copy suppression.
+
+Keep this quest until those cases pass too. Recommend coordinating PSI pacing with media admission and decode deadlines, while retaining every distinct tune-in point and constant-rate slot length. Adding padding after a full slot or deleting distinct tune-in copies would break the existing contracts. Alternatively, explicitly narrow this quest to adjacent-copy suppression and plan the broader pacing work separately; that scope change remains a maintainer decision.

rs/justfile

@@ -352,3 +352,8 @@ bench-rtmp *args:
 [positional-arguments]
 bench-stats *args:
     cargo bench --locked -p moq-stats --features bench --bench producer_tick --bench decode -- "$@"
+
+# TS layout cost across video PIDs and packets per PID.
+[positional-arguments]
+bench-ts-layout *args:
+    cargo bench --locked -p moq-mux --bench ts_layout -- "$@"

rs/moq-mux/Cargo.toml

@@ -50,3 +50,7 @@ tracing-test = { workspace = true, features = ["no-env-filter"] }
 [[bench]]
 name = "catalog"
 harness = false
+
+[[bench]]
+name = "ts_layout"
+harness = false

rs/moq-mux/benches/ts_layout.rs

@@ -0,0 +1,52 @@
+//! TS slot layout across video PID count and packets per PID.
+
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use std::hint::black_box;
+use std::time::Duration;
+
+// Keep the implementation private; layout itself has no codec or transport dependency.
+mod export {
+	pub const PCR_INTERVAL: std::time::Duration = std::time::Duration::from_millis(25);
+}
+#[path = "../src/container/ts/schedule.rs"]
+#[allow(dead_code)]
+mod schedule;
+
+fn packet(pid: u16) -> [u8; 188] {
+	let mut packet = [0xff; 188];
+	packet[..4].copy_from_slice(&[0x47, (pid >> 8) as u8, pid as u8, 0x10]);
+	packet
+}
+
+fn layout(c: &mut Criterion) {
+	let mut group = c.benchmark_group("ts_layout");
+	group.measurement_time(Duration::from_secs(1));
+	group.warm_up_time(Duration::from_millis(200));
+	group.sample_size(10);
+	for videos in [1u16, 4, 16, 64] {
+		for count in [1, 16, 128] {
+			let mut packets = Vec::new();
+			for pid in 1..=videos {
+				packets.extend(packet(0));
+				packets.extend(packet(100));
+				for _ in 0..count {
+					packets.extend(packet(pid));
+				}
+			}
+			let slot = schedule::Slot {
+				index: 0,
+				pcr: 0,
+				packets,
+				nulls: Some(128),
+				keyframe: true,
+				units: (1..=videos).collect(),
+			};
+			group.bench_with_input(BenchmarkId::new(videos.to_string(), count), &slot, |b, slot| {
+				b.iter(|| black_box(slot.layout(black_box(&packet(1)), 100, black_box(&packet(0x1fff)))));
+			});
+		}
+	}
+	group.finish();
+}
+criterion_group!(benches, layout);
+criterion_main!(benches);

rs/moq-mux/src/container/ts/schedule.rs

@@ -136,8 +136,9 @@ pub(super) struct Slot {
 	pub pcr: u64,
 	/// The media packets, in the order they were muxed.
 	pub packets: Vec<u8>,
-	/// How many null packets pad the slot to the multiplex rate, its clock packet included.
-	pub nulls: usize,
+	/// How many null packets pad the slot to the multiplex rate, its clock packet included, or
+	/// `None` when the output is unpadded.
+	pub nulls: Option<usize>,
 	/// Whether a keyframe's first packet rides in the slot.
 	pub keyframe: bool,
 	/// The PID of each access unit whose first packet rides in the slot.
@@ -155,37 +156,59 @@ impl Slot {
 	/// frame's few packets muxed whole, or a video PID taking most of the slot. The program
 	/// tables (PAT and the PMT on `pmt_pid`) go just ahead of the first packet muxed after
 	/// them, so they still lead the keyframe they were written for and a reader knows each PID
-	/// before its first packet.
+	/// before its first packet. Where that would lay a run of them right after an identical
+	/// one, as keyframes aligned across video PIDs write, the copy is dropped (and padded back
+	/// at the end of a padded slot): laid together, the copies would overrun those PIDs'
+	/// transport buffers and mark no tune-in point the first does not.
 	pub fn layout(&self, clock: &[u8], pmt_pid: u16, null: &[u8]) -> Vec<u8> {
 		let packets = self.packets.as_chunks::<{ TsPacket::SIZE }>().0;
-		let mut tables = Vec::new();
-		// Each PID's packets in mux order: the clock's PID first, then the rest as they appear.
+		// The runs of tables in mux order, and each PID's packets: the clock's PID first, then
+		// the rest as they appear.
+		let mut runs: Vec<std::ops::Range<usize>> = Vec::new();
 		let mut lanes: Vec<(u16, VecDeque<usize>)> = vec![(pid(clock), VecDeque::new())];
 		for (i, packet) in packets.iter().enumerate() {
 			let pid = pid(packet);
 			if pid == 0 || pid == pmt_pid {
-				tables.push(i);
+				match runs.last_mut() {
+					Some(run) if run.end == i => run.end += 1,
+					_ => runs.push(i..i + 1),
+				}
 			} else if let Some((_, lane)) = lanes.iter_mut().find(|(p, _)| *p == pid) {
 				lane.push_back(i);
 			} else {
 				lanes.push((pid, VecDeque::from([i])));
 			}
 		}
+		let mut runs = runs.into_iter().peekable();
+		let mut dropped = 0;
+		// Lay the runs muxed ahead of packet `before`, each unless it repeats the one just laid.
+		let mut tables = |out: &mut Vec<u8>, before: usize| {
+			let mut last = None;
+			while let Some(run) = runs.next_if(|run| run.start < before) {
+				let run = &packets[run];
+				if last == Some(run) {
+					dropped += run.len();
+				} else {
+					out.extend_from_slice(run.as_flattened());
+					last = Some(run);
+				}
+			}
+		};
 
 		// Every position, each lane (the nulls last) gains its weight and the one furthest
 		// ahead takes it, paying back the total: over the slot each takes exactly its weight.
 		// The clock takes the first.
+		let nulls = self.nulls.unwrap_or(0);
 		let mut weights: Vec<i64> = lanes
 			.iter()
 			.map(|(_, lane)| lane.len() as i64)
-			.chain([self.nulls as i64])
+			.chain([nulls as i64])
 			.collect();
 		weights[0] += 1;
 		let total: i64 = weights.iter().sum();
 		let mut credit = weights.clone();
 		credit[0] -= total;
-		let mut tables = tables.into_iter().peekable();
-		let mut out = Vec::with_capacity((1 + packets.len() + self.nulls) * TsPacket::SIZE);
+		let mut out = Vec::with_capacity((1 + packets.len() + nulls) * TsPacket::SIZE);
 		out.extend_from_slice(clock);
 		for _ in 1..total {
 			for (credit, weight) in credit.iter_mut().zip(&weights) {
@@ -205,13 +228,14 @@ impl Slot {
 			let i = lane
 				.pop_front()
 				.expect("a lane takes no more positions than its weight");
-			while let Some(table) = tables.next_if(|&table| table < i) {
-				out.extend_from_slice(&packets[table]);
-			}
+			tables(&mut out, i);
 			out.extend_from_slice(&packets[i]);
 		}
-		for table in tables {
-			out.extend_from_slice(&packets[table]);
+		tables(&mut out, packets.len());
+		if self.nulls.is_some() {
+			for _ in 0..dropped {
+				out.extend_from_slice(null);
+			}
 		}
 		out
 	}
@@ -402,7 +426,7 @@ impl Schedule {
 					);
 				}
 				let media: usize = take.iter().sum();
-				(take, allowed as usize - 1 - media, rate_pcr(index, rate))
+				(take, Some(allowed as usize - 1 - media), rate_pcr(index, rate))
 			}
 			None => {
 				let mut take = self.admit(index, usize::MAX);
@@ -414,7 +438,7 @@ impl Schedule {
 						*take = unit.remaining();
 					}
 				}
-				(take, 0, grid_pcr(index))
+				(take, None, grid_pcr(index))
 			}
 		};
 
@@ -576,7 +600,7 @@ mod tests {
 			for packet in slot.packets.as_chunks::<{ TsPacket::SIZE }>().0 {
 				*per_pid.entry(pid(packet)).or_default() += 1;
 			}
-			slots.push((slot.index, per_pid, slot.nulls));
+			slots.push((slot.index, per_pid, slot.nulls.unwrap_or(0)));
 		}
 		slots
 	}
@@ -813,7 +837,7 @@ mod tests {
 				"slot {} strays a packet off the grid",
 				slot.index
 			);
-			laid += 1 + (slot.packets.len() / TsPacket::SIZE + slot.nulls) as u128;
+			laid += 1 + (slot.packets.len() / TsPacket::SIZE + slot.nulls.unwrap_or(0)) as u128;
 		}
 	}
 
@@ -826,7 +850,7 @@ mod tests {
 		schedule.push(1, ms(1_100), unit(1, 100), false);
 		// Unpadded, both go out as soon as they are released, a window ahead.
 		let first: Vec<_> = std::iter::from_fn(|| schedule.next(Some(slot(ms(1_100)))).unwrap())
-			.map(|slot| (slot.index, slot.packets.len() / TsPacket::SIZE, slot.nulls))
+			.map(|slot| (slot.index, slot.packets.len() / TsPacket::SIZE, slot.nulls.unwrap_or(0)))
 			.collect();
 		assert_eq!(first, [(36, 8, 0), (37, 0, 0), (38, 0, 0), (39, 0, 0)]);
 		schedule.set_rate(Some(RATE));
@@ -877,7 +901,7 @@ mod tests {
 			index: 0,
 			pcr: 0,
 			packets,
-			nulls: 416 - 1 - media,
+			nulls: Some(416 - 1 - media),
 			keyframe: false,
 			units: vec![],
 		};
@@ -913,7 +937,7 @@ mod tests {
 			index: 0,
 			pcr: 0,
 			packets,
-			nulls: 2,
+			nulls: Some(2),
 			keyframe: true,
 			units: vec![1, 3],
 		};
@@ -922,4 +946,128 @@ mod tests {
 		let laid: Vec<u16> = slot.layout(&unit(1, 1), 100, &null).chunks(188).map(pid).collect();
 		assert_eq!(laid, [1, 0, 100, 3, 0x1fff, 3, 3, 1, 0x1fff, 3]);
 	}
+
+	/// Keyframes aligned across video PIDs each carry the tables ahead of them. A copy laid
+	/// right after the same tables is dropped, and padded back in a padded slot, while a
+	/// changed table and a copy laid in a place of its own still go out.
+	#[test]
+	fn a_repeated_table_run_is_laid_once() {
+		let tables = |version: u8| {
+			let mut run = unit(0, 1);
+			let mut pmt = unit(100, 1);
+			pmt[10] = version;
+			run.extend(pmt);
+			run
+		};
+		// The heavy PID 3 takes the first position after the clock, so the first three runs go
+		// out together; PID 4's packet then separates the last two.
+		let mut packets = Vec::new();
+		for (pid, count, version) in [(1, 1, 0), (2, 1, 0), (3, 20, 0), (4, 1, 1), (5, 1, 1)] {
+			packets.extend(tables(version));
+			packets.extend(unit(pid, count));
+		}
+		let mut null = unit(0, 1);
+		(null[1], null[2]) = (0x1f, 0xff);
+		let lay = |nulls| {
+			let slot = Slot {
+				index: 0,
+				pcr: 0,
+				packets: packets.clone(),
+				nulls,
+				keyframe: true,
+				units: vec![1, 2, 3, 4, 5],
+			};
+			slot.layout(&unit(1, 1), 100, &null)
+				.chunks(188)
+				.map(pid)
+				.collect::<Vec<u16>>()
+		};
+		let at = |laid: &[u16], p: u16| {
+			laid.iter()
+				.enumerate()
+				.filter(|&(_, &q)| q == p)
+				.map(|(k, _)| k)
+				.collect::<Vec<_>>()
+		};
+
+		let padded = lay(Some(2));
+		assert_eq!(padded.len(), 1 + 10 + 24 + 2, "the slot keeps its size");
+		assert_eq!(at(&padded, 0).len(), 3, "one copy of the first three runs");
+		let pmts = at(&padded, 100);
+		assert_eq!(pmts.len(), 3);
+		assert!(pmts[1] < at(&padded, 4)[0] && pmts[2] < at(&padded, 5)[0]);
+		assert_eq!(at(&padded, 0x1fff).len(), 2 + 4, "the dropped copies pad the slot");
+
+		let unpadded = lay(None);
+		assert_eq!(unpadded.len(), 1 + 6 + 24, "an unpadded slot gains no nulls");
+		assert!(at(&unpadded, 0x1fff).is_empty());
+	}
+	/// Aligned keyframes must not burst repeated PSI past either system transport buffer.
+	#[test]
+	fn aligned_keyframes_respect_system_transport_buffers() {
+		for rate in [25_000_000u64, 50_000_000, 100_000_000] {
+			for videos in [3u16, 4, 8] {
+				let mut packets = Vec::new();
+				for video in 1..=videos {
+					packets.extend(unit(0, 1));
+					packets.extend(unit(100, 1));
+					packets.extend(unit(video as u8, if video == videos { 100 } else { 1 }));
+				}
+				let count = packets.len() / TsPacket::SIZE;
+				let slot = Slot {
+					index: 0,
+					pcr: 0,
+					packets,
+					nulls: Some(rate as usize / PACKET as usize - 1 - count),
+					keyframe: true,
+					units: (1..=videos).collect(),
+				};
+				let mut null = unit(0xff, 1);
+				null[1] = 0x1f;
+				let laid = slot.layout(&unit(1, 1), 100, &null);
+				// ISO 13818-1 2.4.2.3: each system PID's 512-byte TB drains at Rxsys.
+				let drain = Buffer::SYSTEM.rate as f64 / rate as f64 * TsPacket::SIZE as f64;
+				for system in [0, 100] {
+					let (mut fill, mut peak) = (0.0f64, 0.0f64);
+					for packet in laid.chunks(TsPacket::SIZE) {
+						fill = (fill - drain).max(0.0);
+						if pid(packet) == system {
+							fill += TsPacket::SIZE as f64;
+							peak = peak.max(fill);
+						}
+					}
+					assert!(
+						peak <= 512.0,
+						"PID {system} TB peak {peak} B, {videos} videos at {rate} b/s"
+					);
+				}
+			}
+		}
+	}
+	#[test]
+	fn changed_tables_in_the_same_flush_stay_in_order() {
+		let mut packets = Vec::new();
+		for (pid, count, version) in [(1, 1, 0), (2, 1, 1), (3, 20, 1)] {
+			packets.extend(unit(0, 1));
+			let mut pmt = unit(100, 1);
+			pmt[10] = version;
+			packets.extend(pmt);
+			packets.extend(unit(pid, count));
+		}
+		let slot = Slot {
+			index: 0,
+			pcr: 0,
+			packets,
+			nulls: None,
+			keyframe: true,
+			units: vec![1, 2, 3],
+		};
+		let laid = slot.layout(&unit(1, 1), 100, &[]);
+		let versions: Vec<_> = laid
+			.chunks(TsPacket::SIZE)
+			.filter(|packet| pid(packet) == 100)
+			.map(|packet| packet[10])
+			.collect();
+		assert_eq!(versions, [0, 1]);
+	}
 }