use std::collections::HashMap;
use std::task::Poll;
use std::time::Duration;
use bytes::Bytes;
use hang::catalog::{Catalog, Container, VideoConfig};
use mp4_atom::{DecodeMaybe, Encode};
use crate::Result;
use crate::catalog::Stream;
use crate::container::ExportSource;
use crate::container::fmp4::Error;
use crate::container::{Frame, Timestamp};
pub struct Export<S: Stream> {
broadcast: moq_net::BroadcastConsumer,
catalog: Option<S>,
latency: Duration,
fragment_duration: Option<Duration>,
tracks: HashMap<String, Fmp4Track>,
catalog_snapshot: Option<Catalog>,
init_emitted: bool,
}
#[derive(Clone, Debug)]
pub struct Fragment {
pub data: Bytes,
pub init: bool,
pub independent: bool,
pub duration: f64,
}
struct Fmp4Track {
source: ExportSource,
pending: Option<Frame>,
buffer: Vec<Frame>,
buffer_independent: bool,
is_video: bool,
default_frame: Duration,
finished: bool,
track_id: u32,
timescale: u64,
sequence_number: u32,
}
impl<S: Stream> Export<S> {
pub fn new(broadcast: moq_net::BroadcastConsumer, catalog: S) -> Self {
Self {
broadcast,
catalog: Some(catalog),
latency: Duration::ZERO,
fragment_duration: None,
tracks: HashMap::new(),
catalog_snapshot: None,
init_emitted: false,
}
}
pub fn with_latency(mut self, latency: Duration) -> Self {
self.latency = latency;
self
}
pub fn with_fragment_duration(mut self, duration: impl Into<Option<Duration>>) -> Self {
self.fragment_duration = duration.into();
self
}
pub async fn next(&mut self) -> Result<Option<Bytes>> {
Ok(self.next_fragment().await?.map(|f| f.data))
}
pub fn poll_next(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Bytes>>> {
Poll::Ready(Ok(std::task::ready!(self.poll_next_fragment(waiter)?).map(|f| f.data)))
}
pub async fn next_fragment(&mut self) -> Result<Option<Fragment>> {
kio::wait(|waiter| self.poll_next_fragment(waiter)).await
}
pub fn poll_next_fragment(&mut self, waiter: &kio::Waiter) -> Poll<Result<Option<Fragment>>> {
while let Some(catalog) = self.catalog.as_mut() {
match catalog.poll_next(waiter)? {
Poll::Ready(Some(snapshot)) => self.update_catalog(&snapshot.media())?,
Poll::Ready(None) => {
self.catalog = None;
break;
}
Poll::Pending => break,
}
}
let waiting_for_init = !self.init_emitted;
for track in self.tracks.values_mut() {
if track.pending.is_some() || track.finished {
continue;
}
loop {
match track.source.poll_read(waiter) {
Poll::Ready(Ok(Some(frame))) => {
if waiting_for_init && !track.source.header_ready() {
continue;
}
track.pending = Some(frame);
break;
}
Poll::Ready(Ok(None)) => {
track.finished = true;
break;
}
Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
Poll::Pending => break,
}
}
}
if !self.init_emitted {
if self.init_ready() {
let init = self.build_init()?;
self.init_emitted = true;
return Poll::Ready(Ok(Some(Fragment {
data: init,
init: true,
independent: false,
duration: 0.0,
})));
}
if self.catalog.is_none() && self.tracks.values().all(|t| t.finished) {
return Poll::Ready(Ok(None));
}
return Poll::Pending;
}
let chosen = self
.tracks
.iter()
.filter_map(|(name, t)| t.pending.as_ref().map(|f| (name.clone(), f.timestamp)))
.min_by_key(|(_, ts)| *ts)
.map(|(name, _)| name);
if let Some(name) = chosen {
let frag = self.fragment_duration;
let has_video_track = self.tracks.values().any(|t| t.is_video);
let track = self.tracks.get_mut(&name).unwrap();
let frame = track.pending.take().unwrap();
let flush_before = should_flush(track, &frame, frag, has_video_track);
if flush_before {
let frames = std::mem::take(&mut track.buffer);
let fragment = emit_fragment(track, frames, Some(&frame))?;
track.buffer_independent = frame.keyframe;
track.buffer.push(frame);
return Poll::Ready(Ok(Some(fragment)));
}
if track.buffer.is_empty() {
track.buffer_independent = frame.keyframe;
}
track.buffer.push(frame);
return self.poll_next_fragment(waiter);
}
let flushable = self
.tracks
.iter()
.filter_map(|(name, t)| {
if t.finished && !t.buffer.is_empty() {
Some((name.clone(), t.buffer.first().unwrap().timestamp))
} else {
None
}
})
.min_by_key(|(_, ts)| *ts)
.map(|(name, _)| name);
if let Some(name) = flushable {
let track = self.tracks.get_mut(&name).unwrap();
let frames = std::mem::take(&mut track.buffer);
let fragment = emit_fragment(track, frames, None)?;
return Poll::Ready(Ok(Some(fragment)));
}
if self.catalog.is_none() && self.tracks.values().all(|t| t.finished && t.buffer.is_empty()) {
return Poll::Ready(Ok(None));
}
self.tracks
.retain(|_, t| !(t.finished && t.pending.is_none() && t.buffer.is_empty()));
Poll::Pending
}
fn update_catalog(&mut self, catalog: &Catalog) -> Result<()> {
let mut active: HashMap<String, ()> = HashMap::new();
for name in catalog.video.renditions.keys() {
active.insert(name.clone(), ());
}
for name in catalog.audio.renditions.keys() {
active.insert(name.clone(), ());
}
let mut next_track_id = self.tracks.values().map(|t| t.track_id).max().unwrap_or(0) + 1;
for (name, config) in &catalog.video.renditions {
if self.tracks.contains_key(name) {
continue;
}
let source = ExportSource::for_video(&self.broadcast, name, config, self.latency)?;
let timescale = catalog_timescale_video(config);
let framerate = config
.framerate
.filter(|fps| fps.is_finite() && *fps > 0.0)
.unwrap_or(30.0);
self.tracks.insert(
name.clone(),
Fmp4Track {
source,
pending: None,
buffer: Vec::new(),
buffer_independent: false,
is_video: true,
default_frame: Duration::from_secs_f64(1.0 / framerate),
finished: false,
track_id: next_track_id,
timescale,
sequence_number: 1,
},
);
next_track_id += 1;
}
for (name, config) in &catalog.audio.renditions {
if self.tracks.contains_key(name) {
continue;
}
let source = ExportSource::for_audio(&self.broadcast, name, config, self.latency)?;
let timescale = catalog_timescale_audio(config);
self.tracks.insert(
name.clone(),
Fmp4Track {
source,
pending: None,
buffer: Vec::new(),
buffer_independent: false,
is_video: false,
default_frame: Duration::from_secs_f64(1024.0 / config.sample_rate.max(1) as f64),
finished: false,
track_id: next_track_id,
timescale,
sequence_number: 1,
},
);
next_track_id += 1;
}
self.tracks.retain(|name, _| active.contains_key(name));
self.catalog_snapshot = Some(catalog.clone());
Ok(())
}
fn init_ready(&self) -> bool {
self.catalog_snapshot.is_some() && self.tracks.values().all(|t| t.source.header_ready())
}
fn build_init(&self) -> Result<Bytes> {
let catalog = self.catalog_snapshot.as_ref().ok_or(Error::NoCatalogSnapshot)?;
let mut traks: Vec<mp4_atom::Trak> = Vec::new();
let mut trexs: Vec<mp4_atom::Trex> = Vec::new();
let mut ftyp_data: Option<mp4_atom::Ftyp> = None;
for (name, config) in &catalog.video.renditions {
let track = self
.tracks
.get(name)
.ok_or_else(|| Error::MissingVideoTrack(name.clone()))?;
match &config.container {
Container::Cmaf { init, .. } => {
extract_init(init, track.track_id, &mut ftyp_data, &mut traks, &mut trexs)?;
}
Container::Legacy | Container::Loc => {
let description = track.source.description();
let trak = crate::container::fmp4::synthesize_video_trak(
track.track_id,
track.timescale,
config,
description.map(|d| d.as_ref()),
)?;
trexs.push(mp4_atom::Trex {
track_id: trak.tkhd.track_id,
default_sample_description_index: 1,
..Default::default()
});
traks.push(trak);
}
}
}
for (name, config) in &catalog.audio.renditions {
let track = self
.tracks
.get(name)
.ok_or_else(|| Error::MissingAudioTrack(name.clone()))?;
match &config.container {
Container::Cmaf { init, .. } => {
extract_init(init, track.track_id, &mut ftyp_data, &mut traks, &mut trexs)?;
}
Container::Legacy | Container::Loc => {
let trak = crate::container::fmp4::synthesize_audio_trak(track.track_id, track.timescale, config)?;
trexs.push(mp4_atom::Trex {
track_id: trak.tkhd.track_id,
default_sample_description_index: 1,
..Default::default()
});
traks.push(trak);
}
}
}
let ftyp = ftyp_data.unwrap_or(mp4_atom::Ftyp {
major_brand: b"isom".into(),
minor_version: 0x200,
compatible_brands: vec![b"isom".into(), b"iso6".into(), b"mp41".into()],
});
let timescale = traks.first().map(|t| t.mdia.mdhd.timescale).unwrap_or(1000);
let moov = mp4_atom::Moov {
mvhd: mp4_atom::Mvhd {
timescale,
..Default::default()
},
trak: traks,
mvex: if trexs.is_empty() {
None
} else {
Some(mp4_atom::Mvex {
trex: trexs,
..Default::default()
})
},
..Default::default()
};
let mut buf = Vec::new();
ftyp.encode(&mut buf)?;
moov.encode(&mut buf)?;
Ok(Bytes::from(buf))
}
}
fn extract_init(
init: &Bytes,
track_id: u32,
ftyp_data: &mut Option<mp4_atom::Ftyp>,
traks: &mut Vec<mp4_atom::Trak>,
trexs: &mut Vec<mp4_atom::Trex>,
) -> Result<()> {
let mut cursor = std::io::Cursor::new(init.as_ref());
while let Some(atom) = mp4_atom::Any::decode_maybe(&mut cursor)? {
match atom {
mp4_atom::Any::Ftyp(f) if ftyp_data.is_none() => {
*ftyp_data = Some(f);
}
mp4_atom::Any::Moov(moov) => {
for mut trak in moov.trak {
trak.tkhd.track_id = track_id;
traks.push(trak);
}
if let Some(mvex) = moov.mvex {
for mut trex in mvex.trex {
trex.track_id = track_id;
trexs.push(trex);
}
}
}
_ => {}
}
}
Ok(())
}
fn should_flush(track: &Fmp4Track, frame: &Frame, fragment_duration: Option<Duration>, has_video_track: bool) -> bool {
if track.buffer.is_empty() {
return false;
}
if track.is_video && frame.keyframe {
return true;
}
match fragment_duration {
Some(d) if d.is_zero() => true,
Some(d) => {
let mut min = Duration::from(frame.timestamp);
let mut max = min;
for f in &track.buffer {
let pts = Duration::from(f.timestamp);
if pts < min {
min = pts;
}
if pts > max {
max = pts;
}
}
max.saturating_sub(min) >= d
}
None => !track.is_video && !has_video_track,
}
}
fn encode_fragment(track: &mut Fmp4Track, frames: Vec<Frame>) -> Result<Bytes> {
if frames.is_empty() {
return Err(Error::NoFrames.into());
}
let seq = track.sequence_number;
track.sequence_number += 1;
Ok(crate::container::fmp4::encode_fragment(
track.track_id,
track.timescale,
seq,
&frames,
)?)
}
fn emit_fragment(track: &mut Fmp4Track, frames: Vec<Frame>, successor: Option<&Frame>) -> Result<Fragment> {
let independent = !track.is_video || track.buffer_independent;
let frames = infer_missing_durations(frames, successor, track.default_frame);
let duration = fragment_seconds(&frames, track.default_frame);
let data = encode_fragment(track, frames)?;
Ok(Fragment {
data,
init: false,
independent,
duration,
})
}
fn fragment_seconds(frames: &[Frame], default_frame: Duration) -> f64 {
if frames.is_empty() {
return 0.0;
}
if frames
.iter()
.all(|f| f.duration.is_some_and(|duration| !duration.is_zero()))
{
return frames
.iter()
.map(|f| Duration::from(f.duration.unwrap()))
.sum::<Duration>()
.as_secs_f64();
}
let mut min = Duration::MAX;
let mut max = Duration::ZERO;
for f in frames {
let pts = Duration::from(f.timestamp);
min = min.min(pts);
max = max.max(pts);
}
((max - min) + default_frame).as_secs_f64()
}
fn infer_missing_durations(mut frames: Vec<Frame>, successor: Option<&Frame>, default_frame: Duration) -> Vec<Frame> {
let infer_from_pts = pts_monotonic(&frames, successor);
let fallback = Timestamp::try_from(default_frame).ok();
for i in 0..frames.len() {
if frames[i].duration.is_some_and(|duration| !duration.is_zero()) {
continue;
}
frames[i].duration = infer_from_pts
.then(|| next_timestamp(&frames, successor, i))
.flatten()
.and_then(|timestamp| timestamp.checked_sub(frames[i].timestamp).ok())
.filter(|duration| !duration.is_zero())
.or(fallback);
}
frames
}
fn pts_monotonic(frames: &[Frame], successor: Option<&Frame>) -> bool {
let frames_monotonic = frames.windows(2).all(|pair| pair[1].timestamp >= pair[0].timestamp);
let successor_monotonic = match (frames.last(), successor) {
(Some(last), Some(successor)) => successor.timestamp >= last.timestamp,
_ => true,
};
frames_monotonic && successor_monotonic
}
fn next_timestamp(frames: &[Frame], successor: Option<&Frame>, index: usize) -> Option<Timestamp> {
frames
.get(index + 1)
.map(|next| next.timestamp)
.or_else(|| successor.map(|next| next.timestamp))
}
fn catalog_timescale_video(config: &VideoConfig) -> u64 {
match &config.container {
Container::Cmaf { init, .. } => {
parse_timescale_from_init(init).unwrap_or_else(|_| crate::container::fmp4::default_video_timescale(config))
}
Container::Loc | Container::Legacy => crate::container::fmp4::default_video_timescale(config),
}
}
fn catalog_timescale_audio(config: &hang::catalog::AudioConfig) -> u64 {
match &config.container {
Container::Cmaf { init, .. } => parse_timescale_from_init(init).unwrap_or(config.sample_rate as u64),
Container::Loc | Container::Legacy => config.sample_rate as u64,
}
}
fn parse_timescale_from_init(init: &[u8]) -> Result<u64> {
let mut cursor = std::io::Cursor::new(init);
while let Some(atom) = mp4_atom::Any::decode_maybe(&mut cursor)? {
if let mp4_atom::Any::Moov(moov) = atom {
let trak = moov.trak.first().ok_or(Error::NoTracks)?;
return Ok(trak.mdia.mdhd.timescale as u64);
}
}
Err(Error::NoMoov.into())
}
#[cfg(test)]
mod tests {
use bytes::Bytes;
use super::*;
fn ts(micros: u64) -> Timestamp {
Timestamp::from_micros(micros).unwrap()
}
fn frame(timestamp_us: u64, duration_us: Option<u64>) -> Frame {
Frame {
timestamp: ts(timestamp_us),
duration: duration_us.map(ts),
payload: Bytes::from_static(&[0xDE, 0xAD]),
keyframe: false,
}
}
#[test]
fn infer_missing_durations_uses_default_for_trailing_sample() {
let frames = infer_missing_durations(
vec![frame(0, Some(0)), frame(41_667, None), frame(83_334, None)],
None,
Duration::from_millis(33),
);
assert_eq!(frames[0].duration, Some(ts(41_667)));
assert_eq!(frames[1].duration, Some(ts(41_667)));
assert_eq!(frames[2].duration, Some(ts(33_000)));
assert_eq!(fragment_seconds(&frames, Duration::from_millis(33)), 0.116334);
}
#[test]
fn infer_missing_duration_uses_default_for_single_frame() {
let frames = infer_missing_durations(vec![frame(83_333, Some(0))], None, Duration::from_millis(40));
assert_eq!(frames[0].duration, Some(ts(40_000)));
assert_eq!(fragment_seconds(&frames, Duration::from_millis(40)), 0.04);
}
#[test]
fn infer_trailing_duration_from_successor_frame() {
let successor = frame(83_334, None);
let frames = infer_missing_durations(vec![frame(41_667, None)], Some(&successor), Duration::from_millis(33));
assert_eq!(frames[0].duration, Some(ts(41_667)));
assert_eq!(fragment_seconds(&frames, Duration::from_millis(33)), 0.041667);
}
#[test]
fn infer_missing_durations_avoids_non_monotonic_pts() {
let successor = frame(66_000, None);
let frames = infer_missing_durations(
vec![frame(0, None), frame(99_000, None), frame(33_000, None)],
Some(&successor),
Duration::from_millis(33),
);
assert_eq!(frames[0].duration, Some(ts(33_000)));
assert_eq!(frames[1].duration, Some(ts(33_000)));
assert_eq!(frames[2].duration, Some(ts(33_000)));
assert_eq!(fragment_seconds(&frames, Duration::from_millis(33)), 0.099);
}
}