Fix bug seeking to clusters without seekhead

1 files changed, 20 insertions, 3 deletions

diff --git a/transcoder/src/fragment.rs b/transcoder/src/fragment.rs
index c94b877..985c1a4 100644
--- a/transcoder/src/fragment.rs
+++ b/transcoder/src/fragment.rs
@@ -14,7 +14,8 @@ use std::fs::File;
 use std::io::{copy, Write as W2};
 use std::process::{Command, Stdio};
 use std::thread::spawn;
-use winter_matroska::{Segment, TrackEntry as MatroskaTrackEntry};
+use winter_matroska::block::Block;
+use winter_matroska::{Cluster, Segment, TrackEntry as MatroskaTrackEntry};
 
 // TODO odd video resolutions can cause errors when transcoding to YUV42{0,2}
 // TODO with an implementation that cant handle it (SVT-AV1 is such an impl).
@@ -23,7 +24,8 @@ pub fn transcode(
     kind: TrackKind,
     input_key: &str,
     output_format: &StreamFormatInfo,
-    input: Segment,
+    mut input: Segment,
+    next_kf: Option<Block>,
 ) -> Result<Segment> {
     let command = transcode_command(
         kind,
@@ -33,6 +35,9 @@ pub fn transcode(
     )
     .unwrap();
 
+    let input_duration = input.info.duration;
+    let had_next_kf = next_kf.is_some();
+
     let output = cache_file("frag-tc", (input_key, &command), |mut output| {
         let _permit = LOCAL_VIDEO_TRANSCODING_TASKS.lock().unwrap();
         info!("encoding with {command:?}");
@@ -50,6 +55,11 @@ pub fn transcode(
             copy(&mut stdout, &mut output).unwrap();
         });
 
+        input.clusters.extend(next_kf.map(|kf| Cluster {
+            simple_blocks: vec![kf],
+            ..Default::default()
+        }));
+
         write_fragment(ContainerFormat::Matroska, &mut stdin, input)?;
         stdin.flush()?;
         drop(stdin);
@@ -64,12 +74,19 @@ pub fn transcode(
         Box::new(File::open(output.abs())?),
     );
 
-    let info = demuxer.info()?;
+    let mut info = demuxer.info()?;
+    info.duration = input_duration;
     let tracks = demuxer.tracks()?;
     let mut clusters = Vec::new();
     while let Some((_, cluster)) = demuxer.read_cluster()? {
         clusters.push(cluster);
     }
+    if had_next_kf {
+        if let Some(c) = clusters.last_mut() {
+            c.simple_blocks.pop().expect("empty last cluster");
+        }
+        clusters.retain(|c| !c.simple_blocks.is_empty() || !c.block_groups.is_empty());
+    }
 
     Ok(Segment {
         info,