Merge pull request #50 from shiguredo/feature/processor-trait-writer

ライターに MediaProcessor トレイトを実装する

Author: sile

src/composer.rs

@@ -9,9 +9,10 @@ use crate::{
     decoder::VideoDecoderOptions,
     encoder::{AudioEncoder, AudioEncoderThread, VideoEncoder, VideoEncoderThread},
     layout::Layout,
-    media::MediaStreamIdGenerator,
+    media::{MediaSample, MediaStreamId, MediaStreamIdGenerator},
     mixer_audio::AudioMixerThread,
     mixer_video::VideoMixerThread,
+    processor::{MediaProcessor, MediaProcessorInput, MediaProcessorOutput},
     source::{AudioSourceThread, VideoSourceThread},
     stats::{ProcessorStats, Seconds, SharedStats},
     types::CodecName,
@@ -70,25 +71,54 @@ impl Composer {
         );
 
         // 映像ミキサーとエンコーダーを準備
-        let encoded_video_rx = self
+        let mut encoded_video_rx = self
             .create_video_mixer_and_encoder(error_flag.clone(), stats.clone(), video_source_rxs)
             .or_fail()?;
 
         // 音声ミキサーとエンコーダーを準備
-        let encoded_audio_rx = self
+        let mut encoded_audio_rx = self
             .create_audio_mixer_and_encoder(error_flag.clone(), stats.clone(), audio_source_rxs)
             .or_fail()?;
 
         // 合成後の映像と音声への MP4 への書き出しを行う（この処理は現在のスレッドで行う）
+        let writer_input_audio_stream_id = MediaStreamId::new(1000); // audio / video で値が異なっていればなんでもいい
+        let writer_input_video_stream_id = MediaStreamId::new(1001);
+
         let mut mp4_writer = Mp4Writer::new(
             out_file_path,
             &self.layout,
-            encoded_audio_rx,
-            encoded_video_rx,
+            self.layout
+                .has_audio()
+                .then_some(writer_input_audio_stream_id),
+            self.layout
+                .has_video()
+                .then_some(writer_input_video_stream_id),
         )
         .or_fail()?;
 
-        while let Some(timestamp) = mp4_writer.poll().or_fail()? {
+        loop {
+            match mp4_writer.process_output().or_fail()? {
+                MediaProcessorOutput::Finished => break,
+                MediaProcessorOutput::Pending { awaiting_stream_id }
+                    if awaiting_stream_id == writer_input_audio_stream_id =>
+                {
+                    let input = MediaProcessorInput {
+                        stream_id: awaiting_stream_id,
+                        sample: encoded_audio_rx.recv().map(MediaSample::audio_data),
+                    };
+                    mp4_writer.process_input(input).or_fail()?;
+                }
+                MediaProcessorOutput::Pending { awaiting_stream_id } => {
+                    let input = MediaProcessorInput {
+                        stream_id: awaiting_stream_id,
+                        sample: encoded_video_rx.recv().map(MediaSample::video_frame),
+                    };
+                    mp4_writer.process_input(input).or_fail()?;
+                }
+                MediaProcessorOutput::Processed { .. } => unreachable!(),
+            }
+
+            let timestamp = mp4_writer.current_duration();
             progress_bar.set_position(timestamp.as_secs());
             if error_flag.get() {
                 // ファイル読み込み、デコード、合成、エンコード、のいずれかで失敗したものがあるとここに来る

src/writer_mp4.rs

@@ -1,8 +1,10 @@
 use std::{
+    collections::VecDeque,
     fs::File,
     io::{BufWriter, Seek, SeekFrom, Write},
     num::NonZeroU32,
     path::Path,
+    sync::Arc,
     time::{Duration, SystemTime},
 };
 
@@ -17,11 +19,13 @@ use shiguredo_mp4::{
 };
 
 use crate::{
-    audio::{AudioData, AudioDataReceiver},
+    audio::AudioData,
     layout::{Layout, Resolution},
+    media::{MediaSample, MediaStreamId},
     mixer_audio::MIXED_AUDIO_DATA_DURATION,
-    stats::{Mp4WriterStats, Seconds},
-    video::{VideoFrame, VideoFrameReceiver},
+    processor::{MediaProcessor, MediaProcessorInput, MediaProcessorOutput, MediaProcessorSpec},
+    stats::{Mp4WriterStats, ProcessorStats, Seconds},
+    video::VideoFrame,
 };
 
 // Hisui では出力 MP4 のタイムスケールはマイクロ秒固定にする
@@ -42,8 +46,10 @@ pub struct Mp4Writer {
     video_chunks: Vec<Chunk>,
     audio_sample_entry: Option<SampleEntry>,
     video_sample_entry: Option<SampleEntry>,
-    input_audio_rx: Option<AudioDataReceiver>,
-    input_video_rx: Option<VideoFrameReceiver>,
+    input_audio_stream_id: Option<MediaStreamId>,
+    input_video_stream_id: Option<MediaStreamId>,
+    input_audio_queue: VecDeque<Arc<AudioData>>,
+    input_video_queue: VecDeque<Arc<VideoFrame>>,
     finalize_time: Mp4FileTime,
     appending_video_chunk: bool,
     stats: Mp4WriterStats,
@@ -54,8 +60,8 @@ impl Mp4Writer {
     pub fn new<P: AsRef<Path>>(
         path: P,
         layout: &Layout,
-        input_audio_rx: AudioDataReceiver,
-        input_video_rx: VideoFrameReceiver,
+        input_audio_stream_id: Option<MediaStreamId>,
+        input_video_stream_id: Option<MediaStreamId>,
     ) -> orfail::Result<Self> {
         let file = std::fs::OpenOptions::new()
             .create(true)
@@ -74,8 +80,10 @@ impl Mp4Writer {
             audio_sample_entry: None,
             video_sample_entry: None,
             finalize_time: Mp4FileTime::from_unix_time(Duration::ZERO),
-            input_audio_rx: layout.has_audio().then_some(input_audio_rx),
-            input_video_rx: layout.has_video().then_some(input_video_rx),
+            input_audio_stream_id,
+            input_video_stream_id,
+            input_audio_queue: VecDeque::new(),
+            input_video_queue: VecDeque::new(),
             appending_video_chunk: true,
             stats: Mp4WriterStats::default(),
         };
@@ -92,33 +100,16 @@ impl Mp4Writer {
         &self.stats
     }
 
-    /// 新しい入力（合成後の映像と音声）を待機して、それの出力ファイルへの書き込みを行う
-    ///
-    /// 結果は現在の書き込み位置を示すタイムスタンプで、全ての書き込みが完了した場合には `Ok(None)` が返される。
-    pub fn poll(&mut self) -> orfail::Result<Option<Duration>> {
-        let (audio, video) = self.peek_input_audio_and_video();
-        let audio_timestamp = audio.map(|x| x.timestamp);
-        let video_timestamp = video.map(|x| x.timestamp);
-
-        let (result, elapsed) = Seconds::elapsed(|| {
-            self.handle_next_audio_and_video(audio_timestamp, video_timestamp)
-                .or_fail()
-        });
-        self.stats.total_processing_seconds.add(elapsed);
-
-        result
-    }
-
     fn handle_next_audio_and_video(
         &mut self,
         audio_timestamp: Option<Duration>,
         video_timestamp: Option<Duration>,
-    ) -> orfail::Result<Option<Duration>> {
-        match (audio_timestamp,video_timestamp){
+    ) -> orfail::Result<bool> {
+        match (audio_timestamp, video_timestamp){
               (None, None) => {
                 // 全部の入力の処理が完了した
                 self.finalize().or_fail()?;
-                return Ok(None);
+                return Ok(false);
             }
             (None, Some(_)) => {
                 // 残りは映像のみ
@@ -149,30 +140,19 @@ impl Mp4Writer {
             }
         }
 
-        // 進捗（現在のタイムスタンプ）を呼び出し元に返す
-        Ok(Some(self.current_duration()))
+        Ok(true)
     }
 
-    fn current_duration(&self) -> Duration {
+    pub fn current_duration(&self) -> Duration {
         self.stats
             .total_audio_track_seconds
             .get_duration()
             .max(self.stats.total_video_track_seconds.get_duration())
     }
 
-    fn peek_input_audio_and_video(&mut self) -> (Option<&AudioData>, Option<&VideoFrame>) {
-        let audio = self.input_audio_rx.as_mut().and_then(|rx| rx.peek());
-        let video = self.input_video_rx.as_mut().and_then(|rx| rx.peek());
-        (audio, video)
-    }
-
     fn append_video_frame(&mut self, new_chunk: bool) -> orfail::Result<()> {
         // 次の入力を取り出す（これは常に成功する）
-        let frame = self
-            .input_video_rx
-            .as_mut()
-            .and_then(|rx| rx.recv())
-            .or_fail()?;
+        let frame = self.input_video_queue.pop_front().or_fail()?;
 
         if self.stats.video_codec.get().is_none()
             && let Some(name) = frame.format.codec_name()
@@ -184,7 +164,7 @@ impl Mp4Writer {
         // サンプルエントリーは最初に一回だけ存在する
         if self.video_sample_entry.is_none() {
             frame.sample_entry.is_some().or_fail()?;
-            self.video_sample_entry = frame.sample_entry;
+            self.video_sample_entry = frame.sample_entry.clone();
         } else {
             frame.sample_entry.is_none().or_fail()?;
         }
@@ -223,11 +203,7 @@ impl Mp4Writer {
 
     fn append_audio_data(&mut self, new_chunk: bool) -> orfail::Result<()> {
         // 次の入力を取り出す（これは常に成功する）
-        let data = self
-            .input_audio_rx
-            .as_mut()
-            .and_then(|rx| rx.recv())
-            .or_fail()?;
+        let data = self.input_audio_queue.pop_front().or_fail()?;
 
         if self.stats.audio_codec.get().is_none()
             && let Some(name) = data.format.codec_name()
@@ -239,7 +215,7 @@ impl Mp4Writer {
         // サンプルエントリーは最初に一回だけ存在する
         if self.audio_sample_entry.is_none() {
             data.sample_entry.is_some().or_fail()?;
-            self.audio_sample_entry = data.sample_entry;
+            self.audio_sample_entry = data.sample_entry.clone();
         } else {
             data.sample_entry.is_none().or_fail()?;
         }
@@ -743,6 +719,75 @@ impl Mp4Writer {
     }
 }
 
+impl MediaProcessor for Mp4Writer {
+    fn spec(&self) -> MediaProcessorSpec {
+        MediaProcessorSpec {
+            input_stream_ids: self
+                .input_audio_stream_id
+                .into_iter()
+                .chain(self.input_video_stream_id)
+                .collect(),
+            output_stream_ids: Vec::new(),
+            stats: ProcessorStats::Mp4Writer(self.stats.clone()),
+        }
+    }
+
+    fn process_input(&mut self, input: MediaProcessorInput) -> orfail::Result<()> {
+        match input.sample {
+            Some(MediaSample::Audio(sample))
+                if Some(input.stream_id) == self.input_audio_stream_id =>
+            {
+                self.input_audio_queue.push_back(sample);
+            }
+            None if Some(input.stream_id) == self.input_audio_stream_id => {
+                self.input_audio_stream_id = None;
+            }
+            Some(MediaSample::Video(sample))
+                if Some(input.stream_id) == self.input_video_stream_id =>
+            {
+                self.input_video_queue.push_back(sample);
+            }
+            None if Some(input.stream_id) == self.input_video_stream_id => {
+                self.input_video_stream_id = None;
+            }
+            _ => return Err(orfail::Failure::new("BUG: unexpected input stream")),
+        }
+        Ok(())
+    }
+
+    fn process_output(&mut self) -> orfail::Result<MediaProcessorOutput> {
+        loop {
+            if let Some(id) = self.input_video_stream_id
+                && self.input_video_queue.is_empty()
+            {
+                return Ok(MediaProcessorOutput::Pending {
+                    awaiting_stream_id: id,
+                });
+            } else if let Some(id) = self.input_audio_stream_id
+                && self.input_audio_queue.is_empty()
+            {
+                return Ok(MediaProcessorOutput::Pending {
+                    awaiting_stream_id: id,
+                });
+            }
+
+            let audio_timestamp = self.input_audio_queue.front().map(|x| x.timestamp);
+            let video_timestamp = self.input_video_queue.front().map(|x| x.timestamp);
+
+            let (result, elapsed) = Seconds::elapsed(|| {
+                self.handle_next_audio_and_video(audio_timestamp, video_timestamp)
+                    .or_fail()
+            });
+
+            self.stats.total_processing_seconds.add(elapsed);
+
+            if !result? {
+                return Ok(MediaProcessorOutput::Finished);
+            }
+        }
+    }
+}
+
 #[derive(Debug)]
 struct Chunk {
     offset: u64,