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