Add duplicate publisher determinism proof

crates/ec-chopper/src/lib.rs

@@ -12,6 +12,7 @@ use ec_core::{
 };
 use ec_ts::{SectionAssembler, TimeSyncEngine, TimeSyncUpdate, TsReader};
 use serde::{Deserialize, Serialize};
+use std::collections::VecDeque;
 use std::fs;
 use std::io::{Read, Write};
 use std::path::{Path, PathBuf};
@@ -299,12 +300,55 @@ pub fn chunk_ts_stream<T: Read>(
     })
 }
 
+pub fn chunk_ts_stream_with_preroll<T: Read>(
+    stream: T,
+    output_dir: &Path,
+    chunk_duration_ms: u64,
+    max_chunks: Option<usize>,
+    preroll_packets: usize,
+) -> Result<TsChunkManifest> {
+    let mut chunks = Vec::new();
+    chunk_ts_stream_live_with_preroll(
+        stream,
+        output_dir,
+        chunk_duration_ms,
+        max_chunks,
+        preroll_packets,
+        |chunk| {
+            chunks.push(chunk);
+            Ok(())
+        },
+    )?;
+    Ok(TsChunkManifest {
+        output_dir: output_dir.to_path_buf(),
+        chunks,
+    })
+}
+
 pub fn chunk_ts_stream_live<T: Read, F: FnMut(TsChunk) -> Result<()>>(
     stream: T,
     output_dir: &Path,
     chunk_duration_ms: u64,
     max_chunks: Option<usize>,
     mut on_chunk: F,
+) -> Result<()> {
+    chunk_ts_stream_live_with_preroll(
+        stream,
+        output_dir,
+        chunk_duration_ms,
+        max_chunks,
+        0,
+        |chunk| on_chunk(chunk),
+    )
+}
+
+pub fn chunk_ts_stream_live_with_preroll<T: Read, F: FnMut(TsChunk) -> Result<()>>(
+    stream: T,
+    output_dir: &Path,
+    chunk_duration_ms: u64,
+    max_chunks: Option<usize>,
+    preroll_packets: usize,
+    mut on_chunk: F,
 ) -> Result<()> {
     fs::create_dir_all(output_dir)
         .with_context(|| format!("failed to create {}", output_dir.display()))?;
@@ -317,6 +361,7 @@ pub fn chunk_ts_stream_live<T: Read, F: FnMut(TsChunk) -> Result<()>>(
     let mut current_file: Option<std::fs::File> = None;
     let mut current_timing: Option<ChunkTiming> = None;
     let mut emitted = 0usize;
+    let mut preroll = VecDeque::<[u8; ec_ts::TS_PACKET_SIZE]>::with_capacity(preroll_packets);
 
     let mut close_and_emit =
         |index: u64, timing: ChunkTiming, file: std::fs::File| -> Result<bool> {
@@ -332,6 +377,7 @@ pub fn chunk_ts_stream_live<T: Read, F: FnMut(TsChunk) -> Result<()>>(
         };
 
     while let Some(packet) = reader.read_packet()? {
+        let packet_bytes = *packet.as_bytes();
         let updates = engine.ingest_packet(&packet, &mut assembler);
         for update in updates {
             if update.discontinuity {
@@ -344,6 +390,7 @@ pub fn chunk_ts_stream_live<T: Read, F: FnMut(TsChunk) -> Result<()>>(
                         return Ok(());
                     }
                 }
+                preroll.clear();
             }
 
             if let Some(index) = update.chunk_index {
@@ -359,8 +406,11 @@ pub fn chunk_ts_stream_live<T: Read, F: FnMut(TsChunk) -> Result<()>>(
                     }
 
                     let path = chunk_path(output_dir, index);
-                    let file = std::fs::File::create(&path)
+                    let mut file = std::fs::File::create(&path)
                         .with_context(|| format!("failed to create {}", path.display()))?;
+                    for bytes in &preroll {
+                        file.write_all(bytes)?;
+                    }
                     current_file = Some(file);
                     current_index = Some(index);
                     current_timing = Some(ChunkTiming {
@@ -381,6 +431,13 @@ pub fn chunk_ts_stream_live<T: Read, F: FnMut(TsChunk) -> Result<()>>(
         if let Some(file) = current_file.as_mut() {
             file.write_all(packet.as_bytes())?;
         }
+
+        if preroll_packets > 0 {
+            preroll.push_back(packet_bytes);
+            while preroll.len() > preroll_packets {
+                preroll.pop_front();
+            }
+        }
     }
 
     if let (Some(index), Some(timing), Some(file)) = (
@@ -388,7 +445,7 @@ pub fn chunk_ts_stream_live<T: Read, F: FnMut(TsChunk) -> Result<()>>(
         current_timing.take(),
         current_file.take(),
     ) {
-        let _ = close_and_emit(index, timing, file);
+        close_and_emit(index, timing, file)?;
     }
 
     Ok(())
@@ -929,6 +986,43 @@ mod tests {
         let _ = fs::remove_dir_all(&dir);
     }
 
+    #[test]
+    fn chunk_ts_stream_with_preroll_prepends_previous_packets() {
+        let chunk_ms = 1000u64;
+        let dir =
+            std::env::temp_dir().join(format!("ec-chopper-chunks-preroll-{}", std::process::id()));
+        let _ = fs::remove_dir_all(&dir);
+        fs::create_dir_all(&dir).unwrap();
+
+        let packet0 = ts_packet_with_pcr(0x0100, 0, 0);
+        let packet1 = ts_packet_with_pcr(0x0100, 1, 27_000_000);
+        let packet2 = ts_packet_with_pcr(0x0100, 2, 54_000_000);
+        let mut bytes = Vec::new();
+        bytes.extend_from_slice(&packet0);
+        bytes.extend_from_slice(&packet1);
+        bytes.extend_from_slice(&packet2);
+
+        let manifest =
+            chunk_ts_stream_with_preroll(Cursor::new(bytes), &dir, chunk_ms, None, 1).unwrap();
+        let indices = manifest.chunks.iter().map(|c| c.index).collect::<Vec<_>>();
+        assert_eq!(indices, vec![0, 1, 2]);
+
+        assert_eq!(
+            fs::read(&manifest.chunks[0].path).unwrap(),
+            packet0.to_vec()
+        );
+        assert_eq!(
+            fs::read(&manifest.chunks[1].path).unwrap(),
+            [packet0, packet1].concat()
+        );
+        assert_eq!(
+            fs::read(&manifest.chunks[2].path).unwrap(),
+            [packet1, packet2].concat()
+        );
+
+        let _ = fs::remove_dir_all(&dir);
+    }
+
     #[test]
     fn hashed_manifest_merkle_root_matches_core() {
         let dir = std::env::temp_dir().join(format!("ec-chopper-merkle-{}", std::process::id()));