1use std::collections::VecDeque;
12use std::fs;
13use std::path::{Path, PathBuf};
14
15use alpha2_ffmpeg_bridge::OutputFile;
16use alpha2_media::{Container, EncodeError};
17use rsmpeg::avcodec::{AVCodecContext, AVCodecParameters, AVPacket};
18use rsmpeg::avutil::AVRational;
19use rsmpeg::ffi;
20
21use crate::audio_input::{
22 AudioStreamConfig, EncoderAudioInput, OpenedAudioEncoder, open_audio_encoder,
23};
24use crate::convert::{encode_adapter_error, encode_adapter_error_at};
25use crate::encode::{
26 AudioCommand, AudioEncodeThread, AudioEncoderConfig, EncodeStageTimings, QueuedPacket,
27 spawn_audio_encoder, timed,
28};
29use crate::partial::PartialOutput;
30
31pub(crate) struct VideoStream {
33 pub(crate) parameters: AVCodecParameters,
35 pub(crate) time_base: AVRational,
37 pub(crate) copied: bool,
40}
41
42pub(crate) enum AudioStream {
44 Encode(AudioStreamConfig),
46 Copy {
48 parameters: AVCodecParameters,
50 time_base: AVRational,
52 },
53}
54
55enum MuxedAudioTrack {
57 Encoded(MuxedAudio),
58 Copied {
59 stream_index: i32,
60 source_time_base: AVRational,
61 stream_time_base: AVRational,
62 },
63}
64
65pub(crate) struct Muxer {
68 final_path: PathBuf,
69 partial_path: PathBuf,
70 _partial_lock: crate::partial::PartialLock,
72 format: Option<OutputFile>,
73 stream_index: i32,
74 stream_time_base: AVRational,
77 complete: bool,
78 audio: Option<MuxedAudioTrack>,
79 audio_stream_index: i32,
80 audio_packets: VecDeque<QueuedPacket>,
83 pending_header: Option<PendingHeader>,
86 parameter_sets: Option<Vec<u8>>,
90 audio_rescale: Option<(AVRational, AVRational)>,
93}
94
95struct MuxedAudio {
97 thread: AudioEncodeThread,
98 input: EncoderAudioInput,
99}
100
101impl Muxer {
102 #[cfg_attr(
114 not(all(feature = "hwencode", any(target_os = "linux", target_os = "macos"))),
115 allow(dead_code)
116 )]
117 pub(crate) fn open(
118 path: &Path,
119 container: Container,
120 video: &AVCodecContext,
121 video_time_base: AVRational,
122 audio: Option<AudioStreamConfig>,
123 ) -> Result<Self, EncodeError> {
124 Self::open_streams(
125 path,
126 container,
127 VideoStream {
128 parameters: video.extract_codecpar(),
129 time_base: video_time_base,
130 copied: false,
131 },
132 audio.map(AudioStream::Encode),
133 )
134 }
135
136 pub(crate) fn open_streams(
143 path: &Path,
144 container: Container,
145 video: VideoStream,
146 audio: Option<AudioStream>,
147 ) -> Result<Self, EncodeError> {
148 let mut unwritten = Unwritten::open(path, container, video, audio)?;
149 unwritten
150 .format
151 .write_header(&mut crate::convert::muxer_header_options(container))
152 .map_err(encode_adapter_error)?;
153 let stream_time_base = unwritten.time_base(unwritten.stream_index);
154 let audio_stream_time_base = unwritten.time_base(unwritten.audio_stream_index);
155 unwritten.start(stream_time_base, audio_stream_time_base, None)
156 }
157
158 #[cfg_attr(
176 not(any(test, all(feature = "hwencode", target_os = "windows"))),
177 allow(dead_code)
178 )]
179 pub(crate) fn open_deferred(
180 path: &Path,
181 container: Container,
182 video: &AVCodecContext,
183 video_time_base: AVRational,
184 audio: Option<AudioStreamConfig>,
185 ) -> Result<Self, EncodeError> {
186 let unwritten = Unwritten::open(
187 path,
188 container,
189 VideoStream {
190 parameters: video.extract_codecpar(),
191 time_base: video_time_base,
192 copied: false,
193 },
194 audio.map(AudioStream::Encode),
195 )?;
196 let audio_time_base = unwritten
200 .opened_audio
201 .as_ref()
202 .map(|opened| opened.encoder_time_base);
203 unwritten.start(
204 video_time_base,
205 audio_time_base.unwrap_or(video_time_base),
206 Some(PendingHeader {
207 audio_time_base,
208 container,
209 }),
210 )
211 }
212
213 fn write_deferred_header(&mut self, packet: &mut AVPacket) -> Result<(), EncodeError> {
216 let Some(pending) = self.pending_header.as_ref() else {
217 return Ok(());
218 };
219 let parameter_sets =
220 alpha2_ffmpeg_bridge::take_packet_side_data(packet, ffi::AV_PKT_DATA_NEW_EXTRADATA)
221 .filter(|sets| !sets.is_empty())
222 .map(|sets| distinct_parameter_sets(&sets))
223 .ok_or_else(|| {
224 EncodeError::Adapter(
225 "the encoder's first packet carried no parameter sets for the header"
226 .into(),
227 )
228 })?;
229 let (audio_time_base, container) = (pending.audio_time_base, pending.container);
230 let format = self
231 .format
232 .as_mut()
233 .ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?;
234 alpha2_ffmpeg_bridge::set_output_stream_extradata(
235 format,
236 usize::try_from(self.stream_index).unwrap_or(0),
237 ¶meter_sets,
238 )
239 .map_err(|error| encode_adapter_error_at("storing the parameter sets", error))?;
240 crate::encode::beat_encode_call();
241 format
242 .write_header(&mut crate::convert::muxer_header_options(container))
243 .map_err(encode_adapter_error)?;
244 self.stream_time_base =
245 format.streams()[usize::try_from(self.stream_index).unwrap_or(0)].time_base;
246 if let (Some(from), Some(MuxedAudioTrack::Encoded(_))) = (audio_time_base, &self.audio) {
247 let to =
248 format.streams()[usize::try_from(self.audio_stream_index).unwrap_or(0)].time_base;
249 self.audio_rescale = Some((from, to));
250 }
251 self.pending_header = None;
252 self.parameter_sets = Some(parameter_sets);
253 Ok(())
254 }
255
256 pub(crate) fn write_copied_audio(&mut self, packet: &mut AVPacket) -> Result<(), EncodeError> {
263 let Some(MuxedAudioTrack::Copied {
264 stream_index,
265 source_time_base,
266 stream_time_base,
267 }) = self.audio
268 else {
269 return Err(EncodeError::Adapter("this output copies no audio".into()));
270 };
271 packet.set_stream_index(stream_index);
272 packet.rescale_ts(source_time_base, stream_time_base);
273 self.format
274 .as_mut()
275 .ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?
276 .interleaved_write_frame(packet)
277 .map_err(|error| encode_adapter_error_at("packet write failed", error))
278 }
279
280 pub(crate) fn write_video(
292 &mut self,
293 packet: &mut AVPacket,
294 encoder_time_base: AVRational,
295 ) -> Result<(), EncodeError> {
296 if self.pending_header.is_some() {
297 self.write_deferred_header(packet)?;
298 } else if let Some(header) = self.parameter_sets.as_deref() {
299 let repeated =
300 alpha2_ffmpeg_bridge::take_packet_side_data(packet, ffi::AV_PKT_DATA_NEW_EXTRADATA);
301 if repeated.is_some_and(|sets| distinct_parameter_sets(&sets) != header) {
304 return Err(EncodeError::Adapter(
305 "the encoder changed its parameter sets mid-stream, which one stream \
306 header cannot carry"
307 .into(),
308 ));
309 }
310 }
311 packet.set_stream_index(self.stream_index);
312 fill_frame_duration(packet, 1);
314 packet.rescale_ts(encoder_time_base, self.stream_time_base);
315 crate::encode::beat_encode_call();
316 self.format
317 .as_mut()
318 .ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?
319 .interleaved_write_frame(packet)
320 .map_err(|error| encode_adapter_error_at("packet write failed", error))
321 }
322
323 pub(crate) fn write_audio(
330 &mut self,
331 block: &alpha2_media::AudioBlock,
332 stages: &mut EncodeStageTimings,
333 ) -> Result<(), EncodeError> {
334 let Some(MuxedAudioTrack::Encoded(state)) = self.audio.as_mut() else {
335 return Err(EncodeError::Adapter(
336 "this output encodes no audio stream".into(),
337 ));
338 };
339 let frames = state.input.push(block)?;
340 let covered = state.input.received_ticks();
341 for frame in frames {
342 state.thread.send(AudioCommand::Block(frame, covered))?;
343 }
344 self.write_ready_audio(false, stages)
346 }
347
348 pub(crate) fn finish_audio(&mut self) -> Result<(), EncodeError> {
355 let Some(MuxedAudioTrack::Encoded(state)) = self.audio.as_mut() else {
356 return Ok(());
357 };
358 let covered = state.input.received_ticks();
359 for frame in state.input.finish() {
360 state.thread.send(AudioCommand::Block(frame, covered))?;
361 }
362 state.thread.send(AudioCommand::Finish(covered))?;
363 state.thread.shutdown();
364 if let Some(error) = state.thread.lock().failure.clone() {
365 return Err(EncodeError::Adapter(error));
366 }
367 Ok(())
368 }
369
370 pub(crate) fn write_ready_audio(
380 &mut self,
381 flushing: bool,
382 stages: &mut EncodeStageTimings,
383 ) -> Result<(), EncodeError> {
384 let Some(MuxedAudioTrack::Encoded(audio)) = self.audio.as_ref() else {
385 return Ok(());
386 };
387 {
388 let mut state = audio.thread.lock();
389 if let Some(error) = state.failure.clone() {
390 return Err(EncodeError::Adapter(error));
391 }
392 if flushing && !state.finished {
393 return Err(EncodeError::Adapter(
394 "audio encoder stopped before it finished".into(),
395 ));
396 }
397 self.audio_packets.extend(state.packets.drain(..));
398 }
399 if self.pending_header.is_some() {
401 return Ok(());
402 }
403 let rescale = self.audio_rescale;
404 while let Some(mut queued) = self.audio_packets.pop_front() {
405 if let Some((from, to)) = rescale {
406 queued.packet.rescale_ts(from, to);
407 }
408 let format = self
409 .format
410 .as_mut()
411 .ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?;
412 timed(&mut stages.audio_mux_write, || {
413 format.interleaved_write_frame(&mut queued.packet)
414 })
415 .map_err(|error| encode_adapter_error_at("packet write failed", error))?;
416 }
417 Ok(())
418 }
419
420 #[cfg_attr(
423 not(all(
424 feature = "hwencode",
425 any(target_os = "linux", target_os = "macos", target_os = "windows")
426 )),
427 allow(dead_code)
428 )]
429 pub(crate) fn with_audio_timings(&self, stages: EncodeStageTimings) -> EncodeStageTimings {
430 let Some(MuxedAudioTrack::Encoded(audio)) = self.audio.as_ref() else {
431 return stages;
432 };
433 let thread = audio.thread.lock().stages;
434 EncodeStageTimings {
435 audio_send_frame: thread.audio_send_frame,
436 audio_receive_packet: thread.audio_receive_packet,
437 ..stages
438 }
439 }
440
441 pub(crate) fn close(&mut self) -> Result<(), EncodeError> {
448 if self.pending_header.is_some() {
449 return Err(EncodeError::Adapter(
450 "the encoder produced no parameter sets: no video packet reached the header".into(),
451 ));
452 }
453 crate::encode::beat_encode_call();
454 self.format
455 .as_mut()
456 .ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?
457 .write_trailer()
458 .map_err(encode_adapter_error)?;
459 if let Some(output) = self.format.take() {
461 crate::encode::beat_encode_call();
462 output
463 .close()
464 .map_err(|error| encode_adapter_error_at("closing the output", error))?;
465 }
466 crate::partial::publish(&self.partial_path, &self.final_path)?;
467 self.complete = true;
468 Ok(())
469 }
470
471 pub(crate) fn abort(&mut self) {
475 if let Some(MuxedAudioTrack::Encoded(state)) = self.audio.as_mut() {
478 state.thread.shutdown();
479 }
480 if !self.complete {
481 drop(self.format.take());
482 let _ = fs::remove_file(&self.partial_path);
483 }
484 }
485}
486
487impl Drop for Muxer {
488 fn drop(&mut self) {
489 self.abort();
491 }
492}
493
494struct Unwritten {
497 final_path: PathBuf,
498 format: OutputFile,
501 partial: PartialOutput,
502 stream_index: i32,
503 audio_stream_index: i32,
504 opened_audio: Option<OpenedAudioEncoder>,
505 copied_audio: Option<AVRational>,
506}
507
508impl Unwritten {
509 fn open(
515 path: &Path,
516 container: Container,
517 video: VideoStream,
518 audio: Option<AudioStream>,
519 ) -> Result<Self, EncodeError> {
520 let final_path = path.to_owned();
521 let partial = PartialOutput::claim(&final_path)?;
522 let mut format = partial.open_muxer(container)?;
523 let global_header =
524 format.oformat().flags & i32::try_from(ffi::AVFMT_GLOBALHEADER).unwrap_or(0) != 0;
525 if !global_header {
528 return Err(EncodeError::Adapter(format!(
529 "the muxer for {container} has no global header"
530 )));
531 }
532 let stream_index;
533 {
534 let mut stream = format.new_stream();
535 stream.set_time_base(video.time_base);
536 let (codec_id, tag) = (video.parameters.codec_id, video.parameters.codec_tag);
537 stream.set_codecpar(video.parameters);
538 alpha2_ffmpeg_bridge::set_stream_codec_tag(
541 &mut stream,
542 crate::convert::container_codec_tag(container, codec_id, tag, video.copied),
543 );
544 stream_index = stream.index;
545 }
546 let mut opened_audio = None;
547 let mut copied_audio = None;
548 let mut audio_stream_index = 0_i32;
549 match audio {
550 Some(AudioStream::Encode(config)) => {
551 if !container.offers_audio(config.codec.model()) {
552 return Err(EncodeError::Adapter(format!(
553 "{container} does not offer audio codec {}",
554 config.codec.as_str()
555 )));
556 }
557 let opened = open_audio_encoder(config, global_header)?;
558 let mut stream = format.new_stream();
559 stream.set_time_base(opened.encoder_time_base);
560 stream.set_codecpar(opened.encoder.extract_codecpar());
561 audio_stream_index = stream.index;
562 opened_audio = Some(opened);
563 }
564 Some(AudioStream::Copy {
565 parameters,
566 time_base,
567 }) => {
568 let mut stream = format.new_stream();
569 stream.set_time_base(time_base);
570 let (codec_id, tag) = (parameters.codec_id, parameters.codec_tag);
571 stream.set_codecpar(parameters);
572 alpha2_ffmpeg_bridge::set_stream_codec_tag(
573 &mut stream,
574 crate::convert::container_codec_tag(container, codec_id, tag, true),
575 );
576 audio_stream_index = stream.index;
577 copied_audio = Some(time_base);
578 }
579 None => {}
580 }
581 Ok(Self {
582 final_path,
583 format,
584 partial,
585 stream_index,
586 audio_stream_index,
587 opened_audio,
588 copied_audio,
589 })
590 }
591
592 fn time_base(&self, index: i32) -> AVRational {
594 self.format.streams()[usize::try_from(index).unwrap_or(0)].time_base
595 }
596
597 fn start(
603 self,
604 stream_time_base: AVRational,
605 audio_stream_time_base: AVRational,
606 pending_header: Option<PendingHeader>,
607 ) -> Result<Muxer, EncodeError> {
608 let Self {
610 final_path,
611 partial,
612 format,
613 stream_index,
614 audio_stream_index,
615 opened_audio,
616 copied_audio,
617 } = self;
618 let partial_path = partial.path().to_owned();
619 let audio = match (opened_audio, copied_audio) {
622 (Some(opened), _) => Some(MuxedAudioTrack::Encoded(MuxedAudio {
623 input: opened.input,
624 thread: spawn_audio_encoder(AudioEncoderConfig {
625 encoder: opened.encoder,
626 stream_index: audio_stream_index,
627 encoder_time_base: opened.encoder_time_base,
628 stream_time_base: audio_stream_time_base,
629 sample_rate: opened.encoder_rate,
630 sample_format: opened.sample_format,
631 })?,
632 })),
633 (None, Some(source_time_base)) => Some(MuxedAudioTrack::Copied {
634 stream_index: audio_stream_index,
635 source_time_base,
636 stream_time_base: audio_stream_time_base,
637 }),
638 (None, None) => None,
639 };
640 let partial_lock = partial.keep();
641 Ok(Muxer {
642 final_path,
643 partial_path,
644 _partial_lock: partial_lock,
645 format: Some(format),
646 stream_index,
647 stream_time_base,
648 complete: false,
649 audio,
650 audio_stream_index,
651 audio_packets: VecDeque::new(),
652 pending_header,
653 parameter_sets: None,
654 audio_rescale: None,
655 })
656 }
657}
658
659struct PendingHeader {
661 audio_time_base: Option<AVRational>,
664 container: Container,
666}
667
668fn distinct_parameter_sets(sets: &[u8]) -> Vec<u8> {
678 if !(sets.starts_with(&[0, 0, 1]) || sets.starts_with(&[0, 0, 0, 1])) {
679 return sets.to_vec();
680 }
681 let mut starts = Vec::new();
682 let mut at = 0;
683 while at + 3 <= sets.len() {
684 if sets[at..at + 3] == [0, 0, 1] {
685 starts.push(at + 3);
686 at += 3;
687 } else {
688 at += 1;
689 }
690 }
691 let mut units: Vec<&[u8]> = Vec::new();
692 for (index, &start) in starts.iter().enumerate() {
693 let end = starts.get(index + 1).map_or(sets.len(), |next| next - 3);
694 let mut unit = &sets[start..end];
696 while let [rest @ .., 0] = unit {
697 unit = rest;
698 }
699 if !unit.is_empty() && !units.contains(&unit) {
700 units.push(unit);
701 }
702 }
703 units
704 .into_iter()
705 .flat_map(|unit| [&[0, 0, 0, 1][..], unit])
706 .flatten()
707 .copied()
708 .collect()
709}
710
711pub(crate) fn fill_frame_duration(packet: &mut AVPacket, frame_units: i64) {
722 if packet.duration == 0 {
723 packet.set_duration(frame_units);
724 }
725}
726
727pub(crate) fn frame_units(rate: alpha2_time::FrameRate, time_base: AVRational) -> i64 {
730 rsmpeg::avutil::av_rescale_q(
731 1,
732 rsmpeg::avutil::ra(
733 i32::try_from(rate.denominator()).unwrap_or(1),
734 i32::try_from(rate.numerator()).unwrap_or(i32::MAX),
735 ),
736 time_base,
737 )
738 .max(1)
739}