merge keyframe-less clusters with previous when generating cues

2 files changed, 43 insertions, 25 deletions

diff --git a/stream/src/cues.rs b/stream/src/cues.rs
index 2fd4d70..d11e2d7 100644
--- a/stream/src/cues.rs
+++ b/stream/src/cues.rs
@@ -6,7 +6,7 @@
 
 use anyhow::{Result, anyhow};
 use jellycache::{Cache, HashKey};
-use jellyremuxer::demuxers::create_demuxer_autodetect;
+use jellyremuxer::{demuxers::create_demuxer_autodetect, matroska::TrackType};
 use serde::{Deserialize, Serialize};
 use std::{collections::BTreeMap, fs::File, path::Path, sync::Arc};
 
@@ -20,6 +20,7 @@ pub struct TrackStat {
 pub struct GeneratedCue {
     pub position: u64,
     pub time: u64,
+    pub clusters: u64,
 }
 
 #[derive(Serialize, Deserialize)]
@@ -30,28 +31,47 @@ pub struct StatsAndCues {
 
 pub fn generate_cues(cache: &Cache, path: &Path) -> Result<Arc<StatsAndCues>> {
     cache.cache_memory(
-        &format!("media/generated-cues/{}.json", HashKey(path)),
+        &format!("media/generated-cues-v2/{}.json", HashKey(path)),
         move || {
             let media = File::open(path)?;
             let mut media = create_demuxer_autodetect(Box::new(media))?
                 .ok_or(anyhow!("media format unknown"))?;
 
             let info = media.info()?;
+            let tracks = media.tracks()?.ok_or(anyhow!("missing tracks"))?;
+            let video_track_num = tracks
+                .entries
+                .iter()
+                .find(|t| matches!(t.track_type, TrackType::Video))
+                .map(|t| t.track_number);
             media.seek_cluster(None)?;
 
             let mut stats = BTreeMap::<u64, TrackStat>::new();
             let mut cues = Vec::new();
 
             while let Some((position, cluster)) = media.read_cluster()? {
-                cues.push(GeneratedCue {
-                    position,
-                    time: cluster.timestamp * info.timestamp_scale,
-                });
+                let mut has_kf = false;
                 for block in cluster.simple_blocks {
+                    if let Some(vn) = video_track_num {
+                        has_kf |= (block.track == vn) && block.flags.keyframe();
+                    } else {
+                        has_kf = block.flags.keyframe();
+                    }
                     let e = stats.entry(block.track).or_default();
                     e.num_blocks += 1;
                     e.total_size += block.data.len() as u64;
                 }
+                if has_kf {
+                    cues.push(GeneratedCue {
+                        position,
+                        time: cluster.timestamp * info.timestamp_scale,
+                        clusters: 1,
+                    });
+                } else {
+                    cues.last_mut()
+                        .ok_or(anyhow!("first cluster has no keyframe"))?
+                        .clusters += 1;
+                }
             }
 
             Ok(StatsAndCues { stats, cues })
diff --git a/stream/src/fragment.rs b/stream/src/fragment.rs
index 1774e0a..c4de91f 100644
--- a/stream/src/fragment.rs
+++ b/stream/src/fragment.rs
@@ -108,31 +108,29 @@ pub fn fragment_remux(
         .ok_or(anyhow!("fragment index out of range"))?;
     let cluster_offset = start_cue.position;
 
-    let (mut cluster, next_cluster) = {
+    let mut clusters = {
         let media_file = File::open(media_path)?;
         let mut media = create_demuxer_autodetect(Box::new(media_file))?
             .ok_or(anyhow!("media container unknown"))?;
         media.seek_cluster(Some(cluster_offset))?;
-        (
-            media
-                .read_cluster()?
-                .ok_or(anyhow!("cluster unexpectedly missing"))?
-                .1,
-            next_kf
-                .then(|| Ok(media.read_cluster()?.map(|(_, x)| x)))
-                .transpose()?
-                .flatten(),
-        )
+        let mut clusters = Vec::new();
+        for _ in 0..start_cue.clusters {
+            clusters.push(
+                media
+                    .read_cluster()?
+                    .ok_or(anyhow!("cluster unexpectedly missing"))?
+                    .1,
+            );
+        }
+        if next_kf {
+            clusters.extend(media.read_cluster()?.map(|(_, x)| x));
+        }
+        clusters
     };
 
-    cluster.simple_blocks.retain(|b| b.track == file_track_num);
-    cluster
-        .block_groups
-        .retain(|b| b.block.track == file_track_num);
-
-    let mut clusters = vec![cluster];
-    if let Some(next_cluster) = next_cluster {
-        clusters.push(next_cluster);
+    for c in &mut clusters {
+        c.simple_blocks.retain(|b| b.track == file_track_num);
+        c.block_groups.retain(|b| b.block.track == file_track_num);
     }
 
     Ok(Segment {