1use std::collections::VecDeque;
4use std::fs;
5use std::io::{BufWriter, Seek, SeekFrom, Write};
6use std::path::{Path, PathBuf};
7use std::sync::mpsc::{self, Receiver, SyncSender};
8use std::sync::{Arc, Mutex};
9use std::thread::{self, JoinHandle};
10use std::time::{Duration, Instant};
11
12use crate::aac::{AacInput, encoder_sample_rate};
13use crate::capabilities::{
14 AvEncodeConfig, NativeVideoEncoderSelection, VideoEncoderOpenRequest, open_video_encoder,
15};
16use crate::convert::{
17 codec_thread_budget, configure_rgba_to_yuv, encode_adapter_error, encode_adapter_error_at,
18 tick_base,
19};
20use crate::partial::PartialOutput;
21use alpha2_ffmpeg_bridge::OutputFile;
22use alpha2_ffmpeg_bridge::Yuv420pFrameBuffer;
23use alpha2_media::{
24 AudioBlock, AudioBlockSink, CancelToken, DecodedVideoFrame, EncodeError, EncodeEvent,
25 EncodeSink, MAX_AUDIO_CHANNELS, PixelFormat, SpeakerLayout,
26};
27use alpha2_time::MediaTime;
28use rsmpeg::avcodec::{AVCodec, AVCodecContext, AVPacket};
29use rsmpeg::avutil::{AVChannelLayout, AVFrameWithImage, AVImage, AVRational, av_rescale_q, ra};
30use rsmpeg::error::RsmpegError;
31use rsmpeg::ffi;
32use rsmpeg::swscale::SwsContext;
33
34#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
41pub struct EncodeStageTimings {
42 pub video_input_ready_wait: Duration,
46 pub video_send_frame: Duration,
48 pub video_receive_packet: Duration,
50 pub video_mux_write: Duration,
52 pub audio_send_frame: Duration,
54 pub audio_receive_packet: Duration,
56 pub audio_mux_write: Duration,
58}
59
60impl EncodeStageTimings {
61 #[must_use]
63 pub fn since(self, earlier: Self) -> Self {
64 Self {
65 video_input_ready_wait: self
66 .video_input_ready_wait
67 .saturating_sub(earlier.video_input_ready_wait),
68 video_send_frame: self
69 .video_send_frame
70 .saturating_sub(earlier.video_send_frame),
71 video_receive_packet: self
72 .video_receive_packet
73 .saturating_sub(earlier.video_receive_packet),
74 video_mux_write: self.video_mux_write.saturating_sub(earlier.video_mux_write),
75 audio_send_frame: self
76 .audio_send_frame
77 .saturating_sub(earlier.audio_send_frame),
78 audio_receive_packet: self
79 .audio_receive_packet
80 .saturating_sub(earlier.audio_receive_packet),
81 audio_mux_write: self.audio_mux_write.saturating_sub(earlier.audio_mux_write),
82 }
83 }
84}
85
86thread_local! {
87 static ENCODE_CALL_WATCH: std::cell::RefCell<Option<Box<dyn Fn()>>> =
90 const { std::cell::RefCell::new(None) };
91}
92
93#[must_use = "the watch ends when the guard drops"]
104pub fn watch_encode_calls(beat: impl Fn() + 'static) -> EncodeCallWatch {
105 let previous = ENCODE_CALL_WATCH.with(|watch| watch.replace(Some(Box::new(beat))));
106 EncodeCallWatch {
107 previous,
108 thread: std::marker::PhantomData,
109 }
110}
111
112pub struct EncodeCallWatch {
114 previous: Option<Box<dyn Fn()>>,
115 thread: std::marker::PhantomData<*const ()>,
117}
118
119impl Drop for EncodeCallWatch {
120 fn drop(&mut self) {
121 let previous = self.previous.take();
122 ENCODE_CALL_WATCH.with(|watch| {
123 watch.replace(previous);
124 });
125 }
126}
127
128pub(crate) fn beat_encode_call() {
130 ENCODE_CALL_WATCH.with(|watch| {
131 if let Some(beat) = watch.borrow().as_ref() {
132 beat();
133 }
134 });
135}
136
137pub(crate) fn timed<T>(total: &mut Duration, body: impl FnOnce() -> T) -> T {
140 beat_encode_call();
141 let started = Instant::now();
142 let value = body();
143 *total = total.saturating_add(started.elapsed());
144 value
145}
146
147pub struct NativeWavEncodeSink {
149 pub(crate) final_path: PathBuf,
150 pub(crate) partial_path: PathBuf,
151 pub(crate) _partial_lock: crate::partial::PartialLock,
153 pub(crate) file: Option<BufWriter<fs::File>>,
154 pub(crate) sample_rate: u32,
155 pub(crate) channels: u16,
156 pub(crate) speakers: SpeakerLayout,
158 pub(crate) channel_mask: Option<u32>,
162 pub(crate) plane_order: Vec<usize>,
165 pub(crate) frames_written: u64,
166 pub(crate) complete: bool,
167 pub(crate) interleaved: Vec<u8>,
169}
170
171const WAV_WRITE_BUFFER: usize = 1 << 20;
174
175pub struct NativeAvEncodeSink {
180 pub(crate) final_path: PathBuf,
181 pub(crate) partial_path: PathBuf,
182 pub(crate) _partial_lock: crate::partial::PartialLock,
184 pub(crate) output: Option<OutputFile>,
185 pub(crate) video_encoder: AVCodecContext,
186 pub(crate) video_stream_index: i32,
187 pub(crate) video_encoder_time_base: AVRational,
188 pub(crate) video_frame_units: i64,
190 pub(crate) video_stream_time_base: AVRational,
191 pub(crate) video_scaler: SwsContext,
192 pub(crate) yuv_input: crate::convert::Yuv420pConverter,
194 pub(crate) width: i32,
195 pub(crate) height: i32,
196 pub(crate) last_video_pts: Option<i64>,
197 pub(crate) audio: AudioEncodeThread,
200 pub(crate) audio_input: AacInput,
202 pub(crate) video_packets: VecDeque<QueuedPacket>,
203 pub(crate) last_input_ticks: Option<i64>,
204 pub(crate) complete: bool,
205 pub(crate) video_selection: NativeVideoEncoderSelection,
206 pub(crate) stages: EncodeStageTimings,
207}
208
209pub(crate) enum AudioCommand {
211 Block(Vec<Vec<f32>>, i64),
215 Finish(i64),
218}
219
220#[derive(Default)]
225pub(crate) struct AudioQueueState {
226 pub(crate) packets: VecDeque<QueuedPacket>,
228 pub(crate) consumed_ticks: i64,
236 pub(crate) finished: bool,
238 pub(crate) failure: Option<String>,
240 pub(crate) stages: EncodeStageTimings,
242}
243
244pub(crate) struct AudioEncodeThread {
246 commands: Option<SyncSender<AudioCommand>>,
247 join: Option<JoinHandle<()>>,
248 state: Arc<Mutex<AudioQueueState>>,
249}
250
251impl AudioEncodeThread {
252 pub(crate) fn lock(&self) -> std::sync::MutexGuard<'_, AudioQueueState> {
253 self.state
254 .lock()
255 .unwrap_or_else(std::sync::PoisonError::into_inner)
256 }
257
258 pub(crate) fn send(&self, command: AudioCommand) -> Result<(), EncodeError> {
261 let sent = match self.commands.as_ref() {
262 Some(sender) => sender.send(command).is_ok(),
263 None => return Err(EncodeError::Adapter("audio encoder is closed".into())),
264 };
265 if sent {
266 return Ok(());
267 }
268 let failure = self.lock().failure.clone();
269 Err(EncodeError::Adapter(failure.unwrap_or_else(|| {
270 "audio encoder stopped before accepting a block".into()
271 })))
272 }
273
274 pub(crate) fn shutdown(&mut self) {
276 drop(self.commands.take());
277 if let Some(join) = self.join.take() {
278 let _ = join.join();
279 }
280 }
281}
282
283impl Drop for AudioEncodeThread {
284 fn drop(&mut self) {
285 self.shutdown();
286 }
287}
288
289pub(crate) struct QueuedPacket {
291 pub(crate) sort_ticks: i64,
292 pub(crate) packet: AVPacket,
293}
294
295impl NativeAvEncodeSink {
296 pub fn create(path: impl AsRef<Path>, config: AvEncodeConfig) -> Result<Self, EncodeError> {
303 Self::create_selected(path, config, None, &codec_thread_budget)
304 }
305
306 pub fn create_with_video_encoder(
312 path: impl AsRef<Path>,
313 config: AvEncodeConfig,
314 encoder_name: &str,
315 ) -> Result<Self, EncodeError> {
316 Self::create_selected(path, config, Some(encoder_name), &codec_thread_budget)
317 }
318
319 pub fn create_with_video_encoder_threads(
335 path: impl AsRef<Path>,
336 config: AvEncodeConfig,
337 encoder_name: Option<&str>,
338 software_threads: &dyn Fn() -> i64,
339 ) -> Result<Self, EncodeError> {
340 Self::create_selected(path, config, encoder_name, software_threads)
341 }
342
343 #[allow(clippy::too_many_lines)]
344 pub(crate) fn create_selected(
345 path: impl AsRef<Path>,
346 config: AvEncodeConfig,
347 encoder_name: Option<&str>,
348 software_threads: &dyn Fn() -> i64,
349 ) -> Result<Self, EncodeError> {
350 crate::video_encode::check_video_config(&config.video)?;
351 if config.audio_sample_rate == 0
352 || config.audio_channels == 0
353 || config.audio_channels > 8
354 || config.audio_bit_rate == 0
355 {
356 return Err(EncodeError::Adapter(format!(
357 "the audio output ({} channels at {} Hz, {} bit/s) is refused: 1 to 8 channels, a positive rate and bitrate",
358 config.audio_channels, config.audio_sample_rate, config.audio_bit_rate
359 )));
360 }
361 let final_path = path.as_ref().to_owned();
362 let partial = PartialOutput::claim(&final_path)?;
363 let partial_path = partial.path().to_owned();
364 let mut output = partial.open_muxer()?;
365 let global_header =
366 output.oformat().flags & i32::try_from(ffi::AVFMT_GLOBALHEADER).unwrap_or(0) != 0;
367
368 let width = i32::try_from(config.video.width)
369 .map_err(|_| EncodeError::Adapter("width is out of range".into()))?;
370 let height = i32::try_from(config.video.height)
371 .map_err(|_| EncodeError::Adapter("height is out of range".into()))?;
372 let (video_encoder, video_selection) = open_video_encoder(&VideoEncoderOpenRequest {
373 config: &config.video,
374 requested_name: encoder_name,
375 threads: crate::capabilities::SoftwareEncoderThreads::new(software_threads),
376 global_header,
377 width,
378 height,
379 })?;
380 let video_encoder_time_base = video_encoder.time_base;
382
383 let audio_channels = i32::from(config.audio_channels);
384 let audio_codec =
385 AVCodec::find_encoder(config.audio_codec.codec_id()).ok_or_else(|| {
386 EncodeError::Adapter(format!(
387 "{} encoder is unavailable",
388 config.audio_codec.as_str()
389 ))
390 })?;
391 if audio_codec
392 .sample_fmts()
393 .is_some_and(|formats| !formats.contains(&ffi::AV_SAMPLE_FMT_FLTP))
394 {
395 return Err(EncodeError::Adapter(format!(
396 "the {} encoder does not take planar float samples",
397 config.audio_codec.as_str()
398 )));
399 }
400 let encoder_rate = encoder_sample_rate(&audio_codec, config.audio_sample_rate);
401 let audio_sample_rate = i32::try_from(encoder_rate)
402 .map_err(|_| EncodeError::Adapter("audio sample rate is out of range".into()))?;
403 let audio_encoder_time_base = ra(1, audio_sample_rate);
404 let mut audio_encoder = AVCodecContext::new(&audio_codec);
405 audio_encoder.set_sample_rate(audio_sample_rate);
406 audio_encoder.set_sample_fmt(ffi::AV_SAMPLE_FMT_FLTP);
407 audio_encoder.set_ch_layout(AVChannelLayout::from_nb_channels(audio_channels).into_inner());
408 audio_encoder.set_time_base(audio_encoder_time_base);
409 audio_encoder.set_bit_rate(i64::from(config.audio_bit_rate));
410 if global_header {
411 audio_encoder.set_flags(
412 audio_encoder.flags | i32::try_from(ffi::AV_CODEC_FLAG_GLOBAL_HEADER).unwrap_or(0),
413 );
414 }
415 audio_encoder.open(None).map_err(encode_adapter_error)?;
416 let audio_frame_size = usize::try_from(audio_encoder.frame_size)
417 .ok()
418 .filter(|size| *size > 0)
419 .ok_or_else(|| EncodeError::Adapter("AAC encoder has no fixed frame size".into()))?;
420
421 let video_stream_index;
422 {
423 let mut stream = output.new_stream();
424 stream.set_time_base(video_encoder_time_base);
425 stream.set_codecpar(video_encoder.extract_codecpar());
426 crate::video_encode::describe_output_stream(
427 &mut stream,
428 &video_encoder,
429 config.video.display,
430 )?;
431 video_stream_index = stream.index;
432 }
433 let audio_stream_index;
434 {
435 let mut stream = output.new_stream();
436 stream.set_time_base(audio_encoder_time_base);
437 stream.set_codecpar(audio_encoder.extract_codecpar());
438 audio_stream_index = stream.index;
439 }
440 output
441 .write_header(&mut None)
442 .map_err(encode_adapter_error)?;
443 let video_stream_time_base =
444 output.streams()[usize::try_from(video_stream_index).unwrap_or(0)].time_base;
445 let audio_stream_time_base =
446 output.streams()[usize::try_from(audio_stream_index).unwrap_or(0)].time_base;
447 let mut video_scaler = SwsContext::get_context(
448 width,
449 height,
450 ffi::AV_PIX_FMT_RGBA,
451 width,
452 height,
453 ffi::AV_PIX_FMT_YUV420P,
454 ffi::SWS_BILINEAR,
455 None,
456 None,
457 None,
458 )
459 .ok_or_else(|| EncodeError::Adapter("cannot create RGBA to YUV scaler".into()))?;
460 configure_rgba_to_yuv(&mut video_scaler, config.video.color)
461 .map_err(encode_adapter_error)?;
462 let audio = spawn_audio_encoder(AudioEncoderConfig {
463 encoder: audio_encoder,
464 stream_index: audio_stream_index,
465 encoder_time_base: audio_encoder_time_base,
466 stream_time_base: audio_stream_time_base,
467 sample_rate: encoder_rate,
468 })?;
469 let partial_lock = partial.keep();
470 Ok(Self {
471 final_path,
472 partial_path,
473 _partial_lock: partial_lock,
474 output: Some(output),
475 video_encoder,
476 video_stream_index,
477 video_encoder_time_base,
478 video_frame_units: crate::mux::frame_units(
479 config.video.frame_rate,
480 video_encoder_time_base,
481 ),
482 video_stream_time_base,
483 video_scaler,
484 yuv_input: crate::convert::Yuv420pConverter::new(config.video.color),
485 width,
486 height,
487 last_video_pts: None,
488 audio,
489 audio_input: AacInput::new(
490 config.audio_sample_rate,
491 encoder_rate,
492 usize::from(config.audio_channels),
493 audio_frame_size,
494 ),
495 video_packets: VecDeque::new(),
496 last_input_ticks: None,
497 complete: false,
498 video_selection,
499 stages: EncodeStageTimings::default(),
500 })
501 }
502
503 #[cfg(test)]
505 pub(crate) fn queued_audio_packets(&self) -> usize {
506 self.audio.lock().packets.len()
507 }
508
509 #[must_use]
511 pub fn stage_timings(&self) -> EncodeStageTimings {
512 let audio = self.audio.lock().stages;
513 EncodeStageTimings {
514 audio_send_frame: audio.audio_send_frame,
515 audio_receive_packet: audio.audio_receive_packet,
516 ..self.stages
517 }
518 }
519
520 #[must_use]
522 pub fn video_selection(&self) -> &NativeVideoEncoderSelection {
523 &self.video_selection
524 }
525
526 pub(crate) fn cleanup(&mut self) -> Result<(), EncodeError> {
527 self.audio.shutdown();
530 drop(self.output.take());
531 match fs::remove_file(&self.partial_path) {
532 Ok(()) => Ok(()),
533 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
534 Err(error) => Err(encode_adapter_error(error)),
535 }
536 }
537
538 fn queue_video_packet(&mut self, mut packet: AVPacket) {
540 packet.set_stream_index(self.video_stream_index);
541 crate::mux::fill_frame_duration(&mut packet, self.video_frame_units);
542 packet.rescale_ts(self.video_encoder_time_base, self.video_stream_time_base);
543 let timestamp = if packet.dts == ffi::AV_NOPTS_VALUE {
544 packet.pts
545 } else {
546 packet.dts
547 };
548 self.video_packets.push_back(QueuedPacket {
549 sort_ticks: av_rescale_q(timestamp, self.video_stream_time_base, tick_base()),
550 packet,
551 });
552 }
553
554 fn pop_next_packet(&mut self, flushing: bool) -> Result<Option<QueuedPacket>, EncodeError> {
561 let mut state = self.audio.lock();
562 if let Some(error) = state.failure.clone() {
563 return Err(EncodeError::Adapter(error));
564 }
565 let take_audio = match (self.video_packets.front(), state.packets.front()) {
566 (Some(video), Some(audio)) => audio.sort_ticks < video.sort_ticks,
567 (None, Some(_)) => true,
568 (_, None) => false,
569 };
570 if take_audio {
571 return Ok(state.packets.pop_front());
572 }
573 let Some(video) = self.video_packets.front() else {
574 return Ok(None);
575 };
576 if !flushing && !state.finished && state.consumed_ticks < video.sort_ticks {
577 return Ok(None);
578 }
579 drop(state);
580 Ok(self.video_packets.pop_front())
581 }
582
583 pub(crate) fn write_ready_packets(&mut self, flushing: bool) -> Result<(), EncodeError> {
584 while let Some(mut queued) = self.pop_next_packet(flushing)? {
585 let video = queued.packet.stream_index == self.video_stream_index;
587 let total = if video {
588 &mut self.stages.video_mux_write
589 } else {
590 &mut self.stages.audio_mux_write
591 };
592 let output = self
593 .output
594 .as_mut()
595 .ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?;
596 timed(total, || output.interleaved_write_frame(&mut queued.packet))
597 .map_err(|error| encode_adapter_error_at("MP4 packet write failed", error))?;
598 }
599 Ok(())
600 }
601
602 pub(crate) fn drain_video_packets(&mut self, flushing: bool) -> Result<(), EncodeError> {
603 loop {
604 let received = timed(&mut self.stages.video_receive_packet, || {
605 self.video_encoder.receive_packet()
606 });
607 match received {
608 Ok(packet) => self.queue_video_packet(packet),
609 Err(RsmpegError::EncoderDrainError) if !flushing => return Ok(()),
610 Err(RsmpegError::EncoderFlushedError | RsmpegError::EncoderDrainError)
611 if flushing =>
612 {
613 return Ok(());
614 }
615 Err(error) => {
616 return Err(encode_adapter_error_at(
617 "video encoder packet drain failed",
618 error,
619 ));
620 }
621 }
622 }
623 }
624
625 pub fn write_prepared_video(
634 &mut self,
635 frame: &mut Yuv420pFrameBuffer,
636 pts: MediaTime,
637 ) -> Result<(), EncodeError> {
638 let (width, height) = frame.size();
639 if width != usize::try_from(self.width).unwrap_or(0)
640 || height != usize::try_from(self.height).unwrap_or(0)
641 {
642 return Err(EncodeError::Adapter(
643 "encoder input dimensions do not match its configuration".into(),
644 ));
645 }
646 let pts = av_rescale_q(pts.ticks(), tick_base(), self.video_encoder_time_base);
647 if self.last_video_pts.is_some_and(|previous| pts <= previous) {
648 return Err(EncodeError::InvalidTimestamp);
649 }
650 frame.set_presentation(pts, self.video_encoder_time_base);
651 timed(&mut self.stages.video_send_frame, || {
652 self.video_encoder.send_frame(Some(frame.frame()))
653 })
654 .map_err(|error| {
655 encode_adapter_error_at("video encoder rejected a YUV420P frame", error)
656 })?;
657 self.last_video_pts = Some(pts);
658 self.drain_video_packets(false)?;
659 self.write_ready_packets(false)
660 }
661
662 pub(crate) fn write_video(&mut self, frame: &DecodedVideoFrame) -> Result<(), EncodeError> {
663 if frame.visible_rect.width != u32::try_from(self.width).unwrap_or(0)
664 || frame.visible_rect.height != u32::try_from(self.height).unwrap_or(0)
665 {
666 return Err(EncodeError::Adapter(
667 "encoder input dimensions do not match its configuration".into(),
668 ));
669 }
670 let pts = av_rescale_q(frame.pts.ticks(), tick_base(), self.video_encoder_time_base);
671 if self.last_video_pts.is_some_and(|previous| pts <= previous) {
672 return Err(EncodeError::InvalidTimestamp);
673 }
674 if let Some(mut yuv) = self.yuv_input.convert(frame, self.width, self.height)? {
675 yuv.set_pts(pts);
676 yuv.set_time_base(self.video_encoder_time_base);
677 timed(&mut self.stages.video_send_frame, || {
678 self.video_encoder.send_frame(Some(&yuv))
679 })
680 .map_err(|error| {
681 encode_adapter_error_at("video encoder rejected a YUV420P frame", error)
682 })?;
683 self.last_video_pts = Some(pts);
684 self.drain_video_packets(false)?;
685 return self.write_ready_packets(false);
686 }
687 if frame.pixel_format != PixelFormat::Rgba8 || frame.planes.len() != 1 {
688 return Err(EncodeError::Adapter(
689 "encoder expects YUV planes or one RGBA8 plane".into(),
690 ));
691 }
692 let plane = &frame.planes[0];
693 let expected = usize::try_from(self.width)
694 .ok()
695 .and_then(|width| {
696 usize::try_from(self.height)
697 .ok()
698 .and_then(|height| width.checked_mul(height))
699 })
700 .and_then(|pixels| pixels.checked_mul(4))
701 .ok_or_else(|| EncodeError::Adapter("frame dimensions overflow".into()))?;
702 if plane.bytes.len() != expected
703 || plane.stride != u32::try_from(self.width).unwrap_or(0) * 4
704 {
705 return Err(EncodeError::Adapter(
706 "RGBA plane is not tightly packed".into(),
707 ));
708 }
709 let rgba = alpha2_ffmpeg_bridge::rgba_frame(&plane.bytes, self.width, self.height)
713 .map_err(encode_adapter_error)?;
714 let yuv_image = AVImage::new(ffi::AV_PIX_FMT_YUV420P, self.width, self.height, 1)
715 .ok_or_else(|| EncodeError::Adapter("cannot allocate YUV input".into()))?;
716 let mut yuv = AVFrameWithImage::new(yuv_image);
717 self.video_scaler
718 .scale_frame(&rgba, 0, self.height, &mut yuv)
719 .map_err(encode_adapter_error)?;
720 yuv.set_pts(pts);
721 yuv.set_time_base(self.video_encoder_time_base);
722 timed(&mut self.stages.video_send_frame, || {
723 self.video_encoder.send_frame(Some(&yuv))
724 })
725 .map_err(|error| {
726 encode_adapter_error_at("video encoder rejected a converted YUV420P frame", error)
727 })?;
728 self.last_video_pts = Some(pts);
729 self.drain_video_packets(false)?;
730 self.write_ready_packets(false)
731 }
732
733 pub(crate) fn write_audio(&mut self, block: &AudioBlock) -> Result<(), EncodeError> {
736 let frames = self.audio_input.push(block)?;
737 let covered = self.audio_input.received_ticks();
738 for frame in frames {
739 self.audio.send(AudioCommand::Block(frame, covered))?;
740 }
741 self.write_ready_packets(false)
743 }
744}
745
746pub(crate) struct AudioEncoderConfig {
748 pub(crate) encoder: AVCodecContext,
749 pub(crate) stream_index: i32,
750 pub(crate) encoder_time_base: AVRational,
751 pub(crate) stream_time_base: AVRational,
752 pub(crate) sample_rate: u32,
754}
755
756pub(crate) fn spawn_audio_encoder(
766 config: AudioEncoderConfig,
767) -> Result<AudioEncodeThread, EncodeError> {
768 let (commands, receiver) = mpsc::sync_channel(2);
769 let state = Arc::new(Mutex::new(AudioQueueState::default()));
770 let worker_state = Arc::clone(&state);
771 let run = move || run_audio_encoder(config, &receiver, &worker_state);
772 let builder = thread::Builder::new().name("alpha2-audio-encode".to_owned());
773 #[cfg(test)]
774 let spawned = if REFUSE_AUDIO_ENCODE_THREAD.with(std::cell::Cell::take) {
775 drop(run);
776 Err(std::io::Error::from(std::io::ErrorKind::WouldBlock))
777 } else {
778 builder.spawn(run)
779 };
780 #[cfg(not(test))]
781 let spawned = builder.spawn(run);
782 let join = spawned.map_err(|error| {
783 encode_adapter_error_at("failed to start the audio encode thread", error)
784 })?;
785 Ok(AudioEncodeThread {
786 commands: Some(commands),
787 join: Some(join),
788 state,
789 })
790}
791
792#[cfg(test)]
793thread_local! {
794 pub(crate) static REFUSE_AUDIO_ENCODE_THREAD: std::cell::Cell<bool> =
797 const { std::cell::Cell::new(false) };
798}
799
800fn run_audio_encoder(
802 mut config: AudioEncoderConfig,
803 commands: &Receiver<AudioCommand>,
804 state: &Arc<Mutex<AudioQueueState>>,
805) {
806 let mut encoded_samples = 0_u64;
807 loop {
808 let Ok(command) = commands.recv() else {
809 return;
812 };
813 let (channels, covered, last) = match command {
814 AudioCommand::Block(channels, covered) => (Some(channels), covered, false),
815 AudioCommand::Finish(covered) => (None, covered, true),
816 };
817 let mut stages = EncodeStageTimings::default();
818 let outcome = encode_audio_block(
819 &mut config,
820 channels.as_deref(),
821 &mut encoded_samples,
822 &mut stages,
823 );
824 let mut guard = state
825 .lock()
826 .unwrap_or_else(std::sync::PoisonError::into_inner);
827 guard.stages.audio_send_frame = guard
828 .stages
829 .audio_send_frame
830 .saturating_add(stages.audio_send_frame);
831 guard.stages.audio_receive_packet = guard
832 .stages
833 .audio_receive_packet
834 .saturating_add(stages.audio_receive_packet);
835 match outcome {
836 Ok(packets) => guard.packets.extend(packets),
837 Err(error) => {
838 if guard.failure.is_none() {
839 guard.failure = Some(error.to_string());
840 }
841 guard.finished = true;
842 return;
843 }
844 }
845 guard.consumed_ticks = guard.consumed_ticks.max(covered);
846 if last {
847 guard.finished = true;
848 return;
849 }
850 }
851}
852
853fn encode_audio_block(
857 config: &mut AudioEncoderConfig,
858 channels: Option<&[Vec<f32>]>,
859 encoded_samples: &mut u64,
860 stages: &mut EncodeStageTimings,
861) -> Result<Vec<QueuedPacket>, EncodeError> {
862 let last = channels.is_none();
863 if let Some(channels) = channels {
864 let pts = i64::try_from(*encoded_samples).map_err(|_| EncodeError::InvalidTimestamp)?;
865 let frame = alpha2_ffmpeg_bridge::audio_frame_f32p(
866 channels,
867 i32::try_from(config.sample_rate).unwrap_or(i32::MAX),
868 pts,
869 )
870 .map_err(encode_adapter_error)?;
871 timed(&mut stages.audio_send_frame, || {
872 config.encoder.send_frame(Some(&frame))
873 })
874 .map_err(|error| encode_adapter_error_at("audio encoder rejected a frame", error))?;
875 let samples = channels.first().map_or(0, Vec::len);
876 *encoded_samples =
877 encoded_samples.saturating_add(u64::try_from(samples).unwrap_or(u64::MAX));
878 } else {
879 timed(&mut stages.audio_send_frame, || {
880 config.encoder.send_frame(None)
881 })
882 .map_err(encode_adapter_error)?;
883 }
884 let mut packets = Vec::new();
885 loop {
886 let received = timed(&mut stages.audio_receive_packet, || {
887 config.encoder.receive_packet()
888 });
889 match received {
890 Ok(mut packet) => {
891 packet.set_stream_index(config.stream_index);
892 packet.rescale_ts(config.encoder_time_base, config.stream_time_base);
893 let timestamp = if packet.dts == ffi::AV_NOPTS_VALUE {
894 packet.pts
895 } else {
896 packet.dts
897 };
898 packets.push(QueuedPacket {
899 sort_ticks: av_rescale_q(timestamp, config.stream_time_base, tick_base()),
900 packet,
901 });
902 }
903 Err(RsmpegError::EncoderDrainError) if !last => return Ok(packets),
904 Err(RsmpegError::EncoderFlushedError | RsmpegError::EncoderDrainError) if last => {
905 return Ok(packets);
906 }
907 Err(error) => {
908 return Err(encode_adapter_error_at(
909 "audio encoder packet drain failed",
910 error,
911 ));
912 }
913 }
914 }
915}
916
917impl EncodeSink for NativeAvEncodeSink {
918 fn write(&mut self, event: EncodeEvent, cancel: &CancelToken) -> Result<(), EncodeError> {
919 if cancel.is_cancelled() {
920 self.cleanup()?;
921 return Err(EncodeError::Cancelled);
922 }
923 let ticks = match &event {
924 EncodeEvent::Video(frame) => frame.pts.ticks(),
925 EncodeEvent::Audio(block) => block.start_time().ticks(),
926 };
927 if self
928 .last_input_ticks
929 .is_some_and(|previous| ticks < previous)
930 {
931 return Err(EncodeError::InvalidTimestamp);
932 }
933 match event {
934 EncodeEvent::Video(frame) => self.write_video(&frame)?,
935 EncodeEvent::Audio(block) => self.write_audio(&block)?,
936 }
937 self.last_input_ticks = Some(ticks);
938 Ok(())
939 }
940 fn finish(&mut self, cancel: &CancelToken) -> Result<(), EncodeError> {
941 if cancel.is_cancelled() {
942 self.cleanup()?;
943 return Err(EncodeError::Cancelled);
944 }
945 let covered = self.audio_input.received_ticks();
948 for frame in self.audio_input.finish() {
949 self.audio.send(AudioCommand::Block(frame, covered))?;
950 }
951 self.audio.send(AudioCommand::Finish(covered))?;
952 self.audio.shutdown();
953 if let Some(error) = self.audio.lock().failure.clone() {
954 return Err(EncodeError::Adapter(error));
955 }
956 timed(&mut self.stages.video_send_frame, || {
957 self.video_encoder.send_frame(None)
958 })
959 .map_err(encode_adapter_error)?;
960 self.drain_video_packets(true)?;
961 self.write_ready_packets(true)?;
962 beat_encode_call();
963 self.output
964 .as_mut()
965 .ok_or_else(|| EncodeError::Adapter("muxer is already closed".into()))?
966 .write_trailer()
967 .map_err(encode_adapter_error)?;
968 if let Some(output) = self.output.take() {
970 beat_encode_call();
971 output
972 .close()
973 .map_err(|error| encode_adapter_error_at("closing the output", error))?;
974 }
975 crate::partial::publish(&self.partial_path, &self.final_path)?;
976 self.complete = true;
977 Ok(())
978 }
979 fn abort(&mut self) -> Result<(), EncodeError> {
980 self.cleanup()
981 }
982}
983
984impl Drop for NativeAvEncodeSink {
985 fn drop(&mut self) {
986 if !self.complete {
987 let _ = self.cleanup();
988 }
989 }
990}
991
992impl NativeWavEncodeSink {
993 pub fn create(
1000 path: impl AsRef<Path>,
1001 sample_rate: u32,
1002 channels: u16,
1003 ) -> Result<Self, EncodeError> {
1004 let speakers = SpeakerLayout::default_for(usize::from(channels))
1005 .ok_or_else(|| EncodeError::Adapter("invalid WAV format".into()))?;
1006 Self::create_with_speakers(path, sample_rate, speakers)
1007 }
1008
1009 pub fn create_with_speakers(
1021 path: impl AsRef<Path>,
1022 sample_rate: u32,
1023 speakers: SpeakerLayout,
1024 ) -> Result<Self, EncodeError> {
1025 let channels = u16::try_from(speakers.channels())
1026 .ok()
1027 .filter(|_| sample_rate != 0 && speakers.channels() <= MAX_AUDIO_CHANNELS)
1028 .ok_or_else(|| EncodeError::Adapter("invalid WAV format".into()))?;
1029 let (channel_mask, plane_order) = wav_channel_layout(&speakers)?;
1030 let final_path = path.as_ref().to_owned();
1031 let partial = PartialOutput::claim(&final_path)?;
1032 let partial_path = partial.path().to_owned();
1033 let mut file = BufWriter::with_capacity(WAV_WRITE_BUFFER, partial.open_file()?);
1034 file.write_all(&wav_header_with_mask(
1035 sample_rate,
1036 channels,
1037 channel_mask,
1038 0,
1039 )?)
1040 .map_err(encode_adapter_error)?;
1041 let partial_lock = partial.keep();
1042 Ok(Self {
1043 final_path,
1044 partial_path,
1045 _partial_lock: partial_lock,
1046 file: Some(file),
1047 sample_rate,
1048 channels,
1049 speakers,
1050 channel_mask,
1051 plane_order,
1052 frames_written: 0,
1053 complete: false,
1054 interleaved: Vec::new(),
1055 })
1056 }
1057
1058 pub(crate) fn write_block(&mut self, block: &AudioBlock) -> Result<(), EncodeError> {
1059 if block.sample_rate != self.sample_rate
1060 || block.channels.len() != usize::from(self.channels)
1061 || block.channels.is_empty()
1062 {
1063 return Err(EncodeError::Adapter("audio block format changed".into()));
1064 }
1065 if block.speakers != self.speakers {
1066 return Err(EncodeError::Adapter("audio block speakers changed".into()));
1067 }
1068 let frames = block.channels[0].len();
1069 if block.channels.iter().any(|channel| channel.len() != frames) {
1070 return Err(EncodeError::Adapter("audio channel lengths differ".into()));
1071 }
1072 if u64::try_from(block.start_sample).ok() != Some(self.frames_written) {
1073 return Err(EncodeError::InvalidTimestamp);
1074 }
1075 let frames_after = self
1077 .frames_written
1078 .saturating_add(u64::try_from(frames).unwrap_or(u64::MAX));
1079 self.data_bytes(frames_after)?;
1080 let file = self
1081 .file
1082 .as_mut()
1083 .ok_or_else(|| EncodeError::Adapter("WAV sink is closed".into()))?;
1084 self.interleaved.clear();
1085 for frame in 0..frames {
1086 for plane in &self.plane_order {
1087 self.interleaved
1088 .extend_from_slice(&block.channels[*plane][frame].to_le_bytes());
1089 }
1090 }
1091 file.write_all(&self.interleaved)
1092 .map_err(encode_adapter_error)?;
1093 self.frames_written = self
1094 .frames_written
1095 .saturating_add(u64::try_from(frames).unwrap_or(u64::MAX));
1096 Ok(())
1097 }
1098
1099 fn data_bytes(&self, frames: u64) -> Result<u32, EncodeError> {
1102 let overhead = if self.channel_mask.is_some() {
1103 WAV_EXTENSIBLE_OVERHEAD
1104 } else {
1105 WAV_PLAIN_OVERHEAD
1106 };
1107 #[cfg(test)]
1108 let limit = RIFF_LIMIT.with(std::cell::Cell::get);
1109 #[cfg(not(test))]
1110 let limit = u64::from(u32::MAX);
1111 frames
1112 .checked_mul(u64::from(self.channels))
1113 .and_then(|samples| samples.checked_mul(4))
1114 .filter(|bytes| bytes.saturating_add(u64::from(overhead)) <= limit)
1115 .and_then(|bytes| u32::try_from(bytes).ok())
1116 .ok_or_else(|| {
1117 EncodeError::Adapter(format!(
1118 "{frames} frames of {} channels are more than the 4 GiB a WAV file holds",
1119 self.channels
1120 ))
1121 })
1122 }
1123
1124 pub(crate) fn commit(&mut self) -> Result<(), EncodeError> {
1125 let data_bytes = self.data_bytes(self.frames_written)?;
1126 let header = wav_header_with_mask(
1127 self.sample_rate,
1128 self.channels,
1129 self.channel_mask,
1130 data_bytes,
1131 )?;
1132 let mut file = self
1133 .file
1134 .take()
1135 .ok_or_else(|| EncodeError::Adapter("WAV sink is closed".into()))?;
1136 file.seek(SeekFrom::Start(0))
1137 .and_then(|_| file.write_all(&header))
1138 .and_then(|()| file.flush())
1139 .map_err(encode_adapter_error)?;
1140 drop(file);
1141 crate::partial::publish(&self.partial_path, &self.final_path)?;
1142 self.complete = true;
1143 Ok(())
1144 }
1145}
1146
1147impl AudioBlockSink for NativeWavEncodeSink {
1148 fn write_audio(&mut self, block: &AudioBlock, cancel: &CancelToken) -> Result<(), EncodeError> {
1149 if cancel.is_cancelled() {
1150 self.abort_audio()?;
1151 return Err(EncodeError::Cancelled);
1152 }
1153 self.write_block(block)
1154 }
1155 fn finish_audio(&mut self, cancel: &CancelToken) -> Result<(), EncodeError> {
1156 if cancel.is_cancelled() {
1157 self.abort_audio()?;
1158 return Err(EncodeError::Cancelled);
1159 }
1160 self.commit()
1161 }
1162 fn abort_audio(&mut self) -> Result<(), EncodeError> {
1163 drop(self.file.take().map(BufWriter::into_parts));
1165 match fs::remove_file(&self.partial_path) {
1166 Ok(()) => Ok(()),
1167 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
1168 Err(error) => Err(encode_adapter_error(error)),
1169 }
1170 }
1171}
1172
1173impl EncodeSink for NativeWavEncodeSink {
1174 fn write(&mut self, event: EncodeEvent, cancel: &CancelToken) -> Result<(), EncodeError> {
1175 match event {
1176 EncodeEvent::Audio(block) => self.write_audio(&block, cancel),
1177 EncodeEvent::Video(_) => {
1178 Err(EncodeError::Adapter("WAV sink accepts audio only".into()))
1179 }
1180 }
1181 }
1182 fn finish(&mut self, cancel: &CancelToken) -> Result<(), EncodeError> {
1183 self.finish_audio(cancel)
1184 }
1185 fn abort(&mut self) -> Result<(), EncodeError> {
1186 self.abort_audio()
1187 }
1188}
1189
1190impl Drop for NativeWavEncodeSink {
1191 fn drop(&mut self) {
1192 if !self.complete {
1193 let _ = self.abort_audio();
1194 }
1195 }
1196}
1197
1198fn wav_channel_layout(speakers: &SpeakerLayout) -> Result<(Option<u32>, Vec<usize>), EncodeError> {
1201 let identity = (0..speakers.channels()).collect::<Vec<_>>();
1202 if SpeakerLayout::default_for(speakers.channels()) == Some(*speakers) {
1203 return Ok((None, identity));
1204 }
1205 let mut mask = 0_u32;
1206 for speaker in speakers.speakers() {
1207 let bit = speaker.ffmpeg_channel();
1208 if bit >= 18 || mask & (1 << bit) != 0 {
1209 return Err(EncodeError::Adapter(format!(
1210 "a WAVE channel mask cannot say the speakers {}",
1211 speakers.describe()
1212 )));
1213 }
1214 mask |= 1 << bit;
1215 }
1216 let mut order = identity;
1217 order.sort_by_key(|plane| speakers.speakers()[*plane].ffmpeg_channel());
1218 Ok((Some(mask), order))
1219}
1220
1221const WAV_PLAIN_OVERHEAD: u32 = 36;
1223const WAV_EXTENSIBLE_OVERHEAD: u32 = 60;
1225
1226#[cfg(test)]
1227thread_local! {
1228 pub(crate) static RIFF_LIMIT: std::cell::Cell<u64> =
1231 const { std::cell::Cell::new(u32::MAX as u64) };
1232}
1233
1234fn riff_size(data_bytes: u32, overhead: u32) -> Result<u32, EncodeError> {
1235 data_bytes
1236 .checked_add(overhead)
1237 .ok_or_else(|| EncodeError::Adapter("WAV exceeds the RIFF size limit".into()))
1238}
1239
1240pub(crate) fn wav_header_with_mask(
1246 sample_rate: u32,
1247 channels: u16,
1248 channel_mask: Option<u32>,
1249 data_bytes: u32,
1250) -> Result<Vec<u8>, EncodeError> {
1251 let Some(mask) = channel_mask else {
1252 return Ok(wav_header(sample_rate, channels, data_bytes)?.to_vec());
1253 };
1254 let plain = wav_header(sample_rate, channels, 0)?;
1255 let mut header = Vec::with_capacity(68);
1256 header.extend_from_slice(b"RIFF");
1257 header.extend_from_slice(&riff_size(data_bytes, WAV_EXTENSIBLE_OVERHEAD)?.to_le_bytes());
1258 header.extend_from_slice(b"WAVEfmt ");
1259 header.extend_from_slice(&40_u32.to_le_bytes());
1260 header.extend_from_slice(&0xfffe_u16.to_le_bytes());
1261 header.extend_from_slice(&plain[22..36]);
1263 header.extend_from_slice(&22_u16.to_le_bytes());
1264 header.extend_from_slice(&32_u16.to_le_bytes());
1265 header.extend_from_slice(&mask.to_le_bytes());
1266 header.extend_from_slice(&[
1268 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x80, 0x00, 0x00, 0xaa, 0x00, 0x38, 0x9b,
1269 0x71,
1270 ]);
1271 header.extend_from_slice(b"data");
1272 header.extend_from_slice(&data_bytes.to_le_bytes());
1273 Ok(header)
1274}
1275
1276pub(crate) fn wav_header(
1281 sample_rate: u32,
1282 channels: u16,
1283 data_bytes: u32,
1284) -> Result<[u8; 44], EncodeError> {
1285 let mut header = [0_u8; 44];
1286 header[0..4].copy_from_slice(b"RIFF");
1287 header[4..8].copy_from_slice(&riff_size(data_bytes, WAV_PLAIN_OVERHEAD)?.to_le_bytes());
1288 header[8..12].copy_from_slice(b"WAVE");
1289 header[12..16].copy_from_slice(b"fmt ");
1290 header[16..20].copy_from_slice(&16_u32.to_le_bytes());
1291 header[20..22].copy_from_slice(&3_u16.to_le_bytes());
1292 header[22..24].copy_from_slice(&channels.to_le_bytes());
1293 header[24..28].copy_from_slice(&sample_rate.to_le_bytes());
1294 let byte_rate = sample_rate
1295 .saturating_mul(u32::from(channels))
1296 .saturating_mul(4);
1297 header[28..32].copy_from_slice(&byte_rate.to_le_bytes());
1298 header[32..34].copy_from_slice(&channels.saturating_mul(4).to_le_bytes());
1299 header[34..36].copy_from_slice(&32_u16.to_le_bytes());
1300 header[36..40].copy_from_slice(b"data");
1301 header[40..44].copy_from_slice(&data_bytes.to_le_bytes());
1302 Ok(header)
1303}