dash

1 files changed, 57 insertions, 23 deletions

diff --git a/transcoder/src/fragment.rs b/transcoder/src/fragment.rs
index 85abbd5..b6c3b2d 100644
--- a/transcoder/src/fragment.rs
+++ b/transcoder/src/fragment.rs
@@ -6,7 +6,7 @@
 use crate::{Config, LOCAL_VIDEO_TRANSCODING_TASKS};
 use anyhow::Result;
 use jellycache::{Cache, HashKey};
-use jellyremuxer::{ContainerFormat, demuxers::create_demuxer, muxers::write_fragment};
+use jellyremuxer::{ContainerFormat, demuxers::create_demuxer, muxers::write_init_frag};
 use jellystream_types::{StreamFormatInfo, TrackKind};
 use log::info;
 use std::{
@@ -15,30 +15,61 @@ use std::{
     process::{Command, Stdio},
     thread::spawn,
 };
-use winter_matroska::{Cluster, Segment, TrackEntry as MatroskaTrackEntry, block::Block};
+use winter_matroska::{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).
 
-pub fn transcode(
+pub fn transcode_init(
     cache: &Cache,
     config: &Config,
     kind: TrackKind,
     input_key: &str,
     output_format: &StreamFormatInfo,
-    mut input: Segment,
-    next_kf: Option<Block>,
 ) -> Result<Segment> {
     let command = transcode_command(
         kind,
-        &input.tracks.as_ref().unwrap().entries[0],
+        &MatroskaTrackEntry::default(),
         output_format,
+        true,
         config,
-    )
-    .unwrap();
+    )?;
+    let output = cache.cache(
+        &format!(
+            "transcode/media-fragment/{input_key}-{}.mkv",
+            HashKey(&command)
+        ),
+        || {
+            info!("init encode with {command:?}");
+            let mut args = command.split(" ");
+            let proc = Command::new(args.next().unwrap())
+                .stdout(Stdio::piped())
+                .args(args)
+                .spawn()?;
+            Ok(proc.wait_with_output()?.exit_ok()?.stdout)
+        },
+    )?;
 
-    let input_duration = input.info.duration;
-    let had_next_kf = next_kf.is_some();
+    let mut demuxer = create_demuxer(ContainerFormat::Matroska, Box::new(Cursor::new(output)));
+    let info = demuxer.info()?;
+    let tracks = demuxer.tracks()?;
+    Ok(Segment {
+        info,
+        tracks,
+        ..Default::default()
+    })
+}
+
+pub fn transcode(
+    cache: &Cache,
+    config: &Config,
+    kind: TrackKind,
+    input_key: &str,
+    output_format: &StreamFormatInfo,
+    input_track: &MatroskaTrackEntry,
+    segment: impl FnOnce() -> Result<Segment>,
+) -> Result<Segment> {
+    let command = transcode_command(kind, &input_track, output_format, false, config).unwrap();
 
     let output = cache.cache(
         &format!(
@@ -47,6 +78,7 @@ pub fn transcode(
         ),
         || {
             let _permit = LOCAL_VIDEO_TRANSCODING_TASKS.lock().unwrap();
+            let input = segment()?;
             info!("encoding with {command:?}");
             let mut args = command.split(" ");
             let mut proc = Command::new(args.next().unwrap())
@@ -59,11 +91,7 @@ pub fn transcode(
             let mut stdout = proc.stdout.take().unwrap();
 
             spawn(move || {
-                input.clusters.extend(next_kf.map(|kf| Cluster {
-                    simple_blocks: vec![kf],
-                    ..Default::default()
-                }));
-                write_fragment(ContainerFormat::Matroska, &mut stdin, input).unwrap(); // TODO
+                write_init_frag(ContainerFormat::Matroska, &mut stdin, input).unwrap();
                 stdin.flush().unwrap();
                 drop(stdin);
             });
@@ -78,18 +106,19 @@ pub fn transcode(
 
     let mut demuxer = create_demuxer(ContainerFormat::Matroska, Box::new(Cursor::new(output)));
 
-    let mut info = demuxer.info()?;
-    info.duration = input_duration;
+    let info = demuxer.info()?;
     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());
+
+    //? Remove extra kf hack
+    if clusters
+        .last()
+        .map_or(false, |c| c.simple_blocks.len() == 1)
+    {
+        clusters.pop();
     }
 
     Ok(Segment {
@@ -104,6 +133,7 @@ fn transcode_command(
     kind: TrackKind,
     orig_metadata: &MatroskaTrackEntry,
     format: &StreamFormatInfo,
+    dummy: bool,
     config: &Config,
 ) -> Result<String> {
     let br = format.bitrate as u64;
@@ -120,7 +150,11 @@ fn transcode_command(
             write!(o, "-afbc rga ")?;
         }
 
-        write!(o, "-f matroska -i pipe:0 -copyts ")?;
+        if dummy {
+            write!(o, "-f lavfi -i testsrc2 -to 1 ")?;
+        } else {
+            write!(o, "-f matroska -i pipe:0 -copyts ")?;
+        }
 
         if config.enable_rkrga {
             write!(o, "-vf scale_rkrga=w={w}:h={h}:format=nv12:afbc=1 ")?;