mod nv12_dma_buf;
use anyhow::bail;
use ffmpeg_the_third as ffmpeg;
use ffmpeg_the_third::packet::Ref;
use ffmpeg_the_third::sys::*;
use libc::EAGAIN;
use log::{error, info};
pub use nv12_dma_buf::Nv12DmaBufFrame;
use std::ffi::CStr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::Receiver;
use std::sync::{Arc, Condvar, Mutex, mpsc};
use std::thread::JoinHandle;
use std::{ptr, slice, thread};
type StreamResult = (
Option<AudioData>,
Receiver<Nv12DmaBufFrame>,
StreamingController,
JoinHandle<anyhow::Result<()>>,
);
#[derive(Debug, Clone)]
pub struct AudioData {
pub sample_rate: i32,
pub channels: i32,
#[allow(unused)]
pub format: AVSampleFormat, pub pcm_data: Vec<u8>,
}
unsafe extern "C" fn get_hw_format(
_ctx: *mut AVCodecContext,
pix_fmts: *const AVPixelFormat,
) -> AVPixelFormat {
unsafe {
let mut p = pix_fmts;
while !p.is_null() && *p != AVPixelFormat::NONE {
if *p == AVPixelFormat::VAAPI {
return AVPixelFormat::VAAPI;
}
p = p.add(1);
}
}
AVPixelFormat::NONE
}
pub struct VaapiDecoder {
input_ctx: ffmpeg::format::context::Input,
video_stream_index: usize,
codec_ctx: *mut AVCodecContext,
hw_device_ctx: *mut AVBufferRef,
audio_stream_index: Option<usize>,
audio_codec_ctx: *mut AVCodecContext,
time_base: AVRational,
current_pts: i64,
pub codec_name: String,
pub width: i32,
pub height: i32,
pub frame_rate: f64,
pub video_length_ms: i64,
pub cached_frame: Option<Nv12DmaBufFrame>,
}
pub type SeekClosure = Box<dyn Fn(i64, bool) -> anyhow::Result<()> + Send>;
pub struct StreamingController {
pub f_seekr: SeekClosure,
}
impl VaapiDecoder {
pub fn new(file_path: &str) -> Result<Self, String> {
ffmpeg::init().map_err(|e| format!("FFmpeg init failed: {:?}", e))?;
let input_ctx =
ffmpeg::format::input(file_path).map_err(|e| format!("Open input failed: {:?}", e))?;
let video_stream = input_ctx
.streams()
.best(ffmpeg::media::Type::Video)
.ok_or_else(|| "Unable to find video stream.".to_string())?;
let video_stream_index = video_stream.index();
let time_base = video_stream.time_base().into();
let codec_parameters = video_stream.parameters();
let duration = input_ctx.duration();
let video_length_ms = if duration < 0 { 0 } else { duration / 1000 };
unsafe {
let codec = avcodec_find_decoder(codec_parameters.id().into());
if codec.is_null() {
return Err("Unable to find decoder.".into());
}
let codec_name = CStr::from_ptr((*codec).name).to_string_lossy().into_owned();
let mut codec_ctx = avcodec_alloc_context3(codec);
if codec_ctx.is_null() {
return Err("Failed to allocate codec context.".into());
}
let ret = avcodec_parameters_to_context(codec_ctx, codec_parameters.as_ptr());
if ret < 0 {
avcodec_free_context(&mut codec_ctx);
return Err(format!("avcodec_parameters_to_context failed: {}", ret));
}
let mut hw_device_ctx: *mut AVBufferRef = ptr::null_mut();
let ret = av_hwdevice_ctx_create(
&mut hw_device_ctx,
AVHWDeviceType::VAAPI,
ptr::null(),
ptr::null_mut(),
0,
);
if ret < 0 {
avcodec_free_context(&mut codec_ctx);
return Err(format!("VAAPI device init failed: {}", ret));
}
let hw_ref = av_buffer_ref(hw_device_ctx);
if hw_ref.is_null() {
avcodec_free_context(&mut codec_ctx);
av_buffer_unref(&mut hw_device_ctx);
return Err("Failed to reference VAAPI device context.".into());
}
(*codec_ctx).hw_device_ctx = hw_ref;
(*codec_ctx).get_format = Some(get_hw_format);
let ret = avcodec_open2(codec_ctx, codec, ptr::null_mut());
if ret < 0 {
avcodec_free_context(&mut codec_ctx);
av_buffer_unref(&mut hw_device_ctx);
return Err(format!("Failed to open hardware decoder: {}", ret));
}
let width = (*codec_ctx).width;
let height = (*codec_ctx).height;
let mut fr = (*codec_ctx).framerate;
if fr.num == 0 || fr.den == 0 {
let raw_stream = (*input_ctx.as_ptr()).streams.add(video_stream_index);
if !raw_stream.is_null() && !(*raw_stream).is_null() {
fr = (**raw_stream).avg_frame_rate;
}
}
let frame_rate = if fr.num > 0 && fr.den > 0 {
fr.num as f64 / fr.den as f64
} else {
0.0
};
let audio_stream = input_ctx.streams().best(ffmpeg::media::Type::Audio);
let mut audio_stream_index = None;
let mut audio_codec_ctx: *mut AVCodecContext = ptr::null_mut();
if let Some(stream) = audio_stream {
audio_stream_index = Some(stream.index());
let a_codec = avcodec_find_decoder(stream.parameters().id().into());
if !a_codec.is_null() {
audio_codec_ctx = avcodec_alloc_context3(a_codec);
avcodec_parameters_to_context(audio_codec_ctx, stream.parameters().as_ptr());
avcodec_open2(audio_codec_ctx, a_codec, ptr::null_mut());
}
}
Ok(Self {
input_ctx,
video_stream_index,
codec_ctx,
hw_device_ctx,
audio_stream_index,
audio_codec_ctx,
time_base,
current_pts: i64::MIN,
codec_name,
width,
height,
frame_rate,
video_length_ms,
cached_frame: None,
})
}
}
pub fn next_frame(&mut self) -> Result<Option<Nv12DmaBufFrame>, String> {
if let Some(frame) = self.cached_frame.take() {
return Ok(Some(frame));
}
let mut packets = self.input_ctx.packets();
loop {
let mut hw_frame = unsafe { av_frame_alloc() };
if hw_frame.is_null() {
return Err("av_frame_alloc failed".into());
}
let ret = unsafe { avcodec_receive_frame(self.codec_ctx, hw_frame) };
if ret == 0 {
self.current_pts = unsafe { (*hw_frame).pts };
let drm_frame_res = self.map_hw_to_drm_prime(hw_frame);
unsafe { av_frame_free(&mut hw_frame) };
return drm_frame_res.map(Some);
}
unsafe { av_frame_free(&mut hw_frame) };
if ret != AVERROR(EAGAIN) && ret != AVERROR_EOF {
return Err(format!("avcodec_receive_frame failed: {}", ret));
}
if ret == AVERROR_EOF {
return Ok(None);
}
match packets.next() {
Some(Ok((stream, packet))) if stream.index() == self.video_stream_index => unsafe {
let ret = avcodec_send_packet(self.codec_ctx, packet.as_ptr());
if ret < 0 && ret != AVERROR(EAGAIN) && ret != AVERROR_EOF {
return Err(format!("avcodec_send_packet failed: {}", ret));
}
},
Some(Ok(_)) => continue,
Some(Err(e)) => return Err(format!("Reading packet failed: {:?}", e)),
None => unsafe {
let ret = avcodec_send_packet(self.codec_ctx, ptr::null());
if ret < 0 && ret != AVERROR(EAGAIN) && ret != AVERROR_EOF {
return Err(format!("Flush packet failed: {}", ret));
}
},
}
}
}
pub fn seek(&mut self, t_ms: i64, absolute: bool) -> anyhow::Result<()> {
let target_ms = if absolute {
t_ms
} else {
if self.current_pts == i64::MIN {
bail!(
"Cannot seek relatively: current position is unknown (no frames decoded yet)."
);
}
let current_ms = (self.current_pts as i128 * self.time_base.num as i128 * 1000
/ self.time_base.den as i128) as i64;
current_ms + t_ms
}
.max(0);
let target_pts = (target_ms as i128 * self.time_base.den as i128
/ (self.time_base.num as i128 * 1000)) as i64;
unsafe {
let fmt_ctx = self.input_ctx.as_mut_ptr();
let ret = av_seek_frame(
fmt_ctx,
self.video_stream_index as i32,
target_pts,
AVSEEK_FLAG_BACKWARD,
);
if ret < 0 {
bail!(format!("av_seek_frame failed with code: {}", ret));
}
avcodec_flush_buffers(self.codec_ctx);
}
info!("Seeked to {} ms (absolute: {})", target_ms, absolute);
Ok(())
}
fn extract_all_audio(&mut self) -> anyhow::Result<Option<AudioData>> {
let audio_idx = match self.audio_stream_index {
Some(idx) => idx as i32,
None => return Ok(None),
};
if self.audio_codec_ctx.is_null() {
return Ok(None);
}
let mut pcm_data = Vec::new();
let mut sample_rate = 0;
let mut channels = 0;
let target_fmt = AVSampleFormat::S32;
unsafe {
av_seek_frame(self.input_ctx.as_mut_ptr(), -1, 0, AVSEEK_FLAG_BACKWARD);
let mut packet = av_packet_alloc();
let mut frame = av_frame_alloc();
let mut swr_ctx: *mut SwrContext = ptr::null_mut();
let mut process_frame = |f: *mut AVFrame| {
if swr_ctx.is_null() {
sample_rate = (*f).sample_rate;
channels = (*f).ch_layout.nb_channels;
let ret = swr_alloc_set_opts2(
&mut swr_ctx,
&(*f).ch_layout,
target_fmt,
sample_rate,
&(*f).ch_layout,
std::mem::transmute::<i32, AVSampleFormat>((*f).format),
sample_rate,
0,
ptr::null_mut(),
);
if ret < 0 {
error!("swr_alloc_set_opts2 failed: {}", ret);
return;
}
if swr_init(swr_ctx) < 0 {
error!("swr_init failed");
return;
}
}
let in_samples = (*f).nb_samples;
let delay = swr_get_delay(swr_ctx, sample_rate as i64);
let out_samples = in_samples + delay as i32;
let mut out_data: *mut u8 = ptr::null_mut();
let mut out_linesize = 0;
let ret = av_samples_alloc(
&mut out_data,
&mut out_linesize,
channels,
out_samples,
target_fmt,
0,
);
if ret < 0 {
error!("av_samples_alloc failed: {}", ret);
return;
}
let in_data = (*f).data.as_ptr() as *const *const u8;
let converted_samples =
swr_convert(swr_ctx, &out_data, out_samples, in_data, in_samples);
if converted_samples > 0 {
let bytes_per_sample = av_get_bytes_per_sample(target_fmt);
let actual_out_bytes =
(converted_samples * channels * bytes_per_sample) as usize;
let slice = slice::from_raw_parts(out_data, actual_out_bytes);
pcm_data.extend_from_slice(slice);
} else if converted_samples < 0 {
error!("swr_convert failed: {}", converted_samples);
}
av_freep(&mut out_data as *mut _ as *mut std::ffi::c_void);
};
loop {
let ret = av_read_frame(self.input_ctx.as_mut_ptr(), packet);
if ret < 0 {
break;
}
if (*packet).stream_index == audio_idx
&& avcodec_send_packet(self.audio_codec_ctx, packet) == 0
{
loop {
let ret = avcodec_receive_frame(self.audio_codec_ctx, frame);
if ret == AVERROR(EAGAIN) || ret == AVERROR_EOF || ret < 0 {
break;
}
process_frame(frame);
}
}
av_packet_unref(packet);
}
avcodec_send_packet(self.audio_codec_ctx, ptr::null());
loop {
let ret = avcodec_receive_frame(self.audio_codec_ctx, frame);
if ret == AVERROR(EAGAIN) || ret == AVERROR_EOF || ret < 0 {
break;
}
process_frame(frame);
}
if !swr_ctx.is_null() {
loop {
let mut out_data: *mut u8 = ptr::null_mut();
let mut out_linesize = 0;
let out_samples = 8192;
av_samples_alloc(
&mut out_data,
&mut out_linesize,
channels,
out_samples,
target_fmt,
0,
);
let converted_samples =
swr_convert(swr_ctx, &out_data, out_samples, ptr::null(), 0);
if converted_samples > 0 {
let bytes_per_sample = av_get_bytes_per_sample(target_fmt);
let actual_out_bytes =
(converted_samples * channels * bytes_per_sample) as usize;
let slice = slice::from_raw_parts(out_data, actual_out_bytes);
pcm_data.extend_from_slice(slice);
av_freep(&mut out_data as *mut _ as *mut std::ffi::c_void);
} else {
av_freep(&mut out_data as *mut _ as *mut std::ffi::c_void);
break;
}
}
swr_free(&mut swr_ctx);
}
av_packet_free(&mut packet);
av_frame_free(&mut frame);
av_seek_frame(self.input_ctx.as_mut_ptr(), -1, 0, AVSEEK_FLAG_BACKWARD);
avcodec_flush_buffers(self.codec_ctx);
}
if pcm_data.is_empty() {
Ok(None)
} else {
Ok(Some(AudioData {
sample_rate,
channels,
format: target_fmt,
pcm_data,
}))
}
}
pub fn stream<F2>(mut self, mut on_eof: F2) -> anyhow::Result<StreamResult>
where
F2: FnMut() + Send + 'static,
{
let audio_data = self.extract_all_audio()?;
let (tx, rx) = mpsc::sync_channel::<Nv12DmaBufFrame>(3);
let shared_state = Arc::new((Mutex::new((self, false)), Condvar::new()));
let shared_seekr = shared_state.clone();
let g: Arc<AtomicU64> = Arc::new(AtomicU64::new(0));
let ctrl_g = g.clone();
let decoder_thread = thread::spawn(move || -> anyhow::Result<()> {
loop {
let mut guard = match shared_state.0.lock() {
Ok(l) => l,
Err(_) => bail!("Mutex poisoned"),
};
while guard.1 {
guard = match shared_state.1.wait(guard) {
Ok(l) => l,
Err(_) => bail!("Condvar wait failed"),
};
}
let f = guard.0.next_frame();
let mut is_eof_now = false;
let mut frame_to_send = None;
match f {
Ok(Some(mut frame)) => {
frame.generation = g.load(Ordering::Relaxed);
frame_to_send = Some(frame);
}
Ok(None) => {
is_eof_now = true;
guard.1 = true;
}
Err(e) => {
error!("Decoder error: {}", e);
break;
}
}
drop(guard);
if is_eof_now {
info!("Decoder hit EOF, entering standby mode.");
on_eof();
}
if let Some(frame) = frame_to_send
&& tx.send(frame).is_err()
{
info!("Renderer disconnected, shutting down decoder thread.");
break;
}
}
Ok(())
});
Ok((
audio_data,
rx,
StreamingController {
f_seekr: Box::new(move |t: i64, absolute: bool| -> anyhow::Result<()> {
let mut guard = match shared_seekr.0.lock() {
Ok(l) => l,
Err(_) => bail!("Failed to acquire lock in f_seekr()"),
};
let v = &mut guard.0;
let target_ms = if absolute {
t
} else {
if v.current_pts == i64::MIN {
bail!("Cannot seek relatively: current position is unknown (no frames decoded yet).");
}
let current_ms = (v.current_pts as i128 * v.time_base.num as i128 * 1000
/ v.time_base.den as i128) as i64;
current_ms + t
}.max(0);
v.seek(t, absolute)?;
v.cached_frame = None;
while let Ok(Some(frame)) = v.next_frame() {
if frame.timestamp_ms >= target_ms {
v.cached_frame = Some(frame);
break;
}
}
let current = ctrl_g.load(Ordering::Relaxed);
ctrl_g.store(current + 1, Ordering::Relaxed);
guard.1 = false;
shared_seekr.1.notify_all();
Ok(())
}),
},
decoder_thread,
))
}
fn map_hw_to_drm_prime(&self, hw_frame: *mut AVFrame) -> Result<Nv12DmaBufFrame, String> {
unsafe {
let mut drm_frame = av_frame_alloc();
if drm_frame.is_null() {
return Err("av_frame_alloc failed".into());
}
(*drm_frame).format = AVPixelFormat::DRM_PRIME.0;
let ret = av_hwframe_map(drm_frame, hw_frame, 1);
if ret < 0 {
av_frame_free(&mut drm_frame);
return Err(format!("E: VAAPI -> DRM: {}", ret));
}
let desc_ptr = (*drm_frame).data[0] as *const AVDRMFrameDescriptor;
if desc_ptr.is_null() {
av_frame_free(&mut drm_frame);
return Err("DRM Frame im empty".into());
}
let desc = &*desc_ptr;
if desc.nb_objects == 0 || desc.nb_layers == 0 {
av_frame_free(&mut drm_frame);
return Err("Invalid DRM Layout".into());
}
let raw_pts = (*hw_frame).pts;
let timestamp_ms = if raw_pts == i64::MIN {
0
} else {
(raw_pts as i128 * self.time_base.num as i128 * 1000 / self.time_base.den as i128)
as i64
};
let drm_obj = desc.objects[0];
let (pitch_y, offset_y, pitch_uv, offset_uv);
if desc.nb_layers == 1 {
let layer = &desc.layers[0];
if layer.nb_planes < 2 {
av_frame_free(&mut drm_frame);
return Err(format!("Layer : {}", layer.nb_planes));
}
pitch_y = layer.planes[0].pitch as u32;
offset_y = layer.planes[0].offset as u32;
pitch_uv = layer.planes[1].pitch as u32;
offset_uv = layer.planes[1].offset as u32;
} else if desc.nb_layers >= 2 {
let layer_y = &desc.layers[0];
let layer_uv = &desc.layers[1];
if layer_y.nb_planes < 1 || layer_uv.nb_planes < 1 {
av_frame_free(&mut drm_frame);
return Err("Layer".into());
}
pitch_y = layer_y.planes[0].pitch as u32;
offset_y = layer_y.planes[0].offset as u32;
pitch_uv = layer_uv.planes[0].pitch as u32;
offset_uv = layer_uv.planes[0].offset as u32;
} else {
av_frame_free(&mut drm_frame);
return Err("Unknown DRM".into());
}
Ok(Nv12DmaBufFrame {
drm_frame,
fd: drm_obj.fd,
width: (*drm_frame).width,
height: (*drm_frame).height,
timestamp_ms,
format_modifier: drm_obj.format_modifier,
pitch_y,
offset_y,
pitch_uv,
offset_uv,
generation: 0,
})
}
}
}
impl Drop for VaapiDecoder {
fn drop(&mut self) {
unsafe {
if !self.codec_ctx.is_null() {
avcodec_free_context(&mut self.codec_ctx);
}
if !self.audio_codec_ctx.is_null() {
avcodec_free_context(&mut self.audio_codec_ctx);
}
if !self.hw_device_ctx.is_null() {
av_buffer_unref(&mut self.hw_device_ctx);
}
}
}
}
unsafe impl Send for VaapiDecoder {}