mirror of
https://github.com/v0l/zap-stream-core.git
synced 2025-06-20 10:47:02 +00:00
Progress audio encoder (FIFO)
This commit is contained in:
276
src/encode/audio.rs
Normal file
276
src/encode/audio.rs
Normal file
@ -0,0 +1,276 @@
|
||||
use std::mem::transmute;
|
||||
use std::ptr;
|
||||
|
||||
use anyhow::Error;
|
||||
use ffmpeg_sys_next::{
|
||||
av_audio_fifo_alloc, av_audio_fifo_free, av_audio_fifo_read, av_audio_fifo_realloc,
|
||||
av_audio_fifo_size, av_audio_fifo_write, av_buffer_ref, av_buffer_unref,
|
||||
AV_CH_LAYOUT_STEREO, av_channel_layout_copy, av_frame_alloc, av_frame_free, av_frame_get_buffer,
|
||||
av_freep, av_get_sample_fmt, av_packet_alloc, av_packet_free,
|
||||
av_packet_rescale_ts, av_samples_alloc_array_and_samples, AVAudioFifo,
|
||||
AVBufferRef, AVChannelLayout, AVChannelLayout__bindgen_ty_1, AVCodec,
|
||||
avcodec_alloc_context3, avcodec_find_encoder, avcodec_free_context, avcodec_open2, avcodec_receive_packet, avcodec_send_frame,
|
||||
AVCodecContext, AVERROR, AVFrame, swr_alloc_set_opts2, swr_convert, swr_free,
|
||||
swr_init, SwrContext,
|
||||
};
|
||||
use ffmpeg_sys_next::AVChannelOrder::AV_CHANNEL_ORDER_NATIVE;
|
||||
use libc::EAGAIN;
|
||||
use tokio::sync::mpsc::UnboundedSender;
|
||||
|
||||
use crate::ipc::Rx;
|
||||
use crate::pipeline::{PipelinePayload, PipelineProcessor};
|
||||
use crate::utils::{audio_variant_id_ref, get_ffmpeg_error_msg, id_ref_to_uuid};
|
||||
use crate::variant::AudioVariant;
|
||||
|
||||
pub struct AudioEncoder<T> {
|
||||
variant: AudioVariant,
|
||||
ctx: *mut AVCodecContext,
|
||||
codec: *const AVCodec,
|
||||
fifo: *mut AVAudioFifo,
|
||||
swr_ctx: *mut SwrContext,
|
||||
chan_in: T,
|
||||
chan_out: UnboundedSender<PipelinePayload>,
|
||||
var_id_ref: *mut AVBufferRef,
|
||||
}
|
||||
|
||||
unsafe impl<T> Send for AudioEncoder<T> {}
|
||||
|
||||
unsafe impl<T> Sync for AudioEncoder<T> {}
|
||||
|
||||
impl<T> Drop for AudioEncoder<T> {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
swr_free(&mut self.swr_ctx);
|
||||
av_audio_fifo_free(self.fifo);
|
||||
avcodec_free_context(&mut self.ctx);
|
||||
av_buffer_unref(&mut self.var_id_ref);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<TRecv> AudioEncoder<TRecv>
|
||||
where
|
||||
TRecv: Rx<PipelinePayload>,
|
||||
{
|
||||
pub fn new(
|
||||
chan_in: TRecv,
|
||||
chan_out: UnboundedSender<PipelinePayload>,
|
||||
variant: AudioVariant,
|
||||
) -> Self {
|
||||
let id_ref = audio_variant_id_ref(&variant);
|
||||
Self {
|
||||
ctx: ptr::null_mut(),
|
||||
codec: ptr::null(),
|
||||
fifo: ptr::null_mut(),
|
||||
swr_ctx: ptr::null_mut(),
|
||||
variant,
|
||||
chan_in,
|
||||
chan_out,
|
||||
var_id_ref: id_ref,
|
||||
}
|
||||
}
|
||||
|
||||
unsafe fn setup_encoder(&mut self, frame: *mut AVFrame) -> Result<(), Error> {
|
||||
if self.ctx == ptr::null_mut() {
|
||||
let codec = self.variant.codec;
|
||||
let encoder = avcodec_find_encoder(transmute(codec as i32));
|
||||
if encoder == ptr::null_mut() {
|
||||
return Err(Error::msg("Encoder not found"));
|
||||
}
|
||||
|
||||
let ctx = avcodec_alloc_context3(encoder);
|
||||
if ctx == ptr::null_mut() {
|
||||
return Err(Error::msg("Failed to allocate encoder context"));
|
||||
}
|
||||
|
||||
(*ctx).time_base = self.variant.time_base();
|
||||
(*ctx).sample_fmt = av_get_sample_fmt(
|
||||
format!("{}\0", self.variant.sample_fmt).as_ptr() as *const libc::c_char,
|
||||
);
|
||||
(*ctx).bit_rate = self.variant.bitrate as i64;
|
||||
(*ctx).sample_rate = self.variant.sample_rate as libc::c_int;
|
||||
(*ctx).ch_layout = AVChannelLayout {
|
||||
order: AV_CHANNEL_ORDER_NATIVE,
|
||||
nb_channels: 2,
|
||||
u: AVChannelLayout__bindgen_ty_1 {
|
||||
mask: AV_CH_LAYOUT_STEREO,
|
||||
},
|
||||
opaque: ptr::null_mut(),
|
||||
};
|
||||
|
||||
// setup audio FIFO
|
||||
let fifo = av_audio_fifo_alloc((*ctx).sample_fmt, 2, 1);
|
||||
if fifo == ptr::null_mut() {
|
||||
return Err(Error::msg("Failed to allocate audio FiFO buffer"));
|
||||
}
|
||||
|
||||
let mut swr_ctx = ptr::null_mut();
|
||||
let ret = swr_alloc_set_opts2(
|
||||
&mut swr_ctx,
|
||||
&(*ctx).ch_layout,
|
||||
(*ctx).sample_fmt,
|
||||
(*ctx).sample_rate,
|
||||
&(*frame).ch_layout,
|
||||
transmute((*frame).format),
|
||||
(*frame).sample_rate,
|
||||
0,
|
||||
ptr::null_mut(),
|
||||
);
|
||||
if ret < 0 {
|
||||
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
|
||||
}
|
||||
|
||||
let ret = swr_init(swr_ctx);
|
||||
if ret < 0 {
|
||||
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
|
||||
}
|
||||
|
||||
let ret = avcodec_open2(ctx, encoder, ptr::null_mut());
|
||||
if ret < 0 {
|
||||
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
|
||||
}
|
||||
|
||||
self.ctx = ctx;
|
||||
self.codec = encoder;
|
||||
self.swr_ctx = swr_ctx;
|
||||
self.fifo = fifo;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Returns true if we should process audio frame from FIFO
|
||||
/// false if nothing to process this frame
|
||||
unsafe fn process_audio_frame(&mut self, frame: *mut AVFrame) -> Result<bool, Error> {
|
||||
let in_samples = (*frame).nb_samples;
|
||||
let mut dst_samples: *mut *mut u8 = ptr::null_mut();
|
||||
let ret = av_samples_alloc_array_and_samples(
|
||||
&mut dst_samples,
|
||||
ptr::null_mut(),
|
||||
2,
|
||||
in_samples,
|
||||
(*self.ctx).sample_fmt,
|
||||
0,
|
||||
);
|
||||
if ret < 0 {
|
||||
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
|
||||
}
|
||||
|
||||
// resample audio
|
||||
let ret = swr_convert(
|
||||
self.swr_ctx,
|
||||
dst_samples,
|
||||
in_samples,
|
||||
(*frame).extended_data as *const *const u8,
|
||||
in_samples,
|
||||
);
|
||||
if ret < 0 {
|
||||
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
|
||||
}
|
||||
|
||||
// push resampled audio into fifo
|
||||
let ret = av_audio_fifo_realloc(self.fifo, av_audio_fifo_size(self.fifo) + in_samples);
|
||||
if ret < 0 {
|
||||
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
|
||||
}
|
||||
if av_audio_fifo_write(
|
||||
self.fifo,
|
||||
dst_samples as *const *mut libc::c_void,
|
||||
in_samples,
|
||||
) < in_samples
|
||||
{
|
||||
return Err(Error::msg("Failed to write samples to FIFO"));
|
||||
}
|
||||
|
||||
if dst_samples != ptr::null_mut() {
|
||||
av_freep(dst_samples.add(0) as *mut libc::c_void);
|
||||
}
|
||||
|
||||
let buffered = av_audio_fifo_size(self.fifo);
|
||||
Ok(buffered >= (*self.ctx).frame_size)
|
||||
}
|
||||
|
||||
unsafe fn get_fifo_frame(&mut self) -> Result<*mut AVFrame, Error> {
|
||||
let mut frame = av_frame_alloc();
|
||||
let frame_size = (*self.ctx).frame_size.min(av_audio_fifo_size(self.fifo));
|
||||
(*frame).nb_samples = frame_size;
|
||||
av_channel_layout_copy(&mut (*frame).ch_layout, &(*self.ctx).ch_layout);
|
||||
(*frame).format = (*self.ctx).sample_fmt as libc::c_int;
|
||||
(*frame).sample_rate = (*self.ctx).sample_rate;
|
||||
|
||||
let ret = av_frame_get_buffer(frame, 0);
|
||||
if ret < 0 {
|
||||
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
|
||||
}
|
||||
|
||||
let ret = av_audio_fifo_read(
|
||||
self.fifo,
|
||||
ptr::addr_of_mut!((*frame).data) as *const *mut libc::c_void,
|
||||
frame_size,
|
||||
);
|
||||
if ret < frame_size {
|
||||
av_frame_free(&mut frame);
|
||||
return Err(Error::msg("Failed to read frame from FIFO"));
|
||||
}
|
||||
|
||||
Ok(frame)
|
||||
}
|
||||
|
||||
unsafe fn process_frame(&mut self, frame: *mut AVFrame) -> Result<(), Error> {
|
||||
let var_id = id_ref_to_uuid((*frame).opaque_ref)?;
|
||||
assert_eq!(var_id, self.variant.id);
|
||||
|
||||
self.setup_encoder(frame)?;
|
||||
|
||||
if !self.process_audio_frame(frame)? {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// read audio from FIFO
|
||||
let frame = self.get_fifo_frame()?;
|
||||
let mut ret = avcodec_send_frame(self.ctx, frame);
|
||||
if ret < 0 && ret != AVERROR(EAGAIN) {
|
||||
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
|
||||
}
|
||||
|
||||
while ret > 0 || ret == AVERROR(EAGAIN) {
|
||||
let mut pkt = av_packet_alloc();
|
||||
ret = avcodec_receive_packet(self.ctx, pkt);
|
||||
if ret < 0 {
|
||||
av_packet_free(&mut pkt);
|
||||
if ret == AVERROR(EAGAIN) {
|
||||
return Ok(());
|
||||
}
|
||||
return Err(Error::msg(get_ffmpeg_error_msg(ret)));
|
||||
}
|
||||
|
||||
(*pkt).time_base = (*self.ctx).time_base;
|
||||
(*pkt).duration = (*frame).duration;
|
||||
av_packet_rescale_ts(pkt, (*frame).time_base, (*self.ctx).time_base);
|
||||
(*pkt).opaque = self.ctx as *mut libc::c_void;
|
||||
(*pkt).opaque_ref = av_buffer_ref(self.var_id_ref);
|
||||
self.chan_out
|
||||
.send(PipelinePayload::AvPacket("Encoder packet".to_owned(), pkt))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<TRecv> PipelineProcessor for AudioEncoder<TRecv>
|
||||
where
|
||||
TRecv: Rx<PipelinePayload>,
|
||||
{
|
||||
fn process(&mut self) -> Result<(), Error> {
|
||||
while let Ok(pkg) = self.chan_in.try_recv_next() {
|
||||
match pkg {
|
||||
PipelinePayload::AvFrame(_, frm, idx) => unsafe {
|
||||
if self.variant.src_index == idx {
|
||||
self.process_frame(frm)?;
|
||||
}
|
||||
},
|
||||
_ => return Err(Error::msg("Payload not supported")),
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user