Skip to content

Decoupled video buffer from video encoder for more extensibility #12

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Adonca2203 merged 1 commit into from
May 6, 2025
Merged

Decoupled video buffer from video encoder for more extensibility #12

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions src/encoders/buffer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ impl VideoFrameData {
/// The buffer is ordered by decoding timestamp (DTS) and maintains complete GOPs (groups of pictures),
/// ensuring that no partial GOPs are kept when trimming for ease of muxing and playback.
#[derive(Clone)]
pub struct VideoBuffer {
pub struct ShadowCaptureVideoBuffer {
frames: BTreeMap<i64, VideoFrameData>,

/// Maximum duration (in seconds) that the buffer should retain.
Expand All @@ -47,7 +47,7 @@ pub struct VideoBuffer {
key_frame_keys: Vec<i64>,
}

impl VideoBuffer {
impl ShadowCaptureVideoBuffer {
/// Creates a new `FrameBuffer` with a specified maximum duration.
///
/// # Arguments
Expand Down
4 changes: 2 additions & 2 deletions src/encoders/buffer_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ fn test_video_buffer_no_trim() {
VideoFrameData::new(vec![3], true, 6),
];

let mut buffer = VideoBuffer::new(10);
let mut buffer = ShadowCaptureVideoBuffer::new(10);

buffer.insert(1, dummy_frames[0].clone());
buffer.insert(2, dummy_frames[1].clone());
Expand All @@ -42,7 +42,7 @@ fn test_video_buffer_no_trim() {

#[test]
fn test_video_buffer_trimming() {
let mut buffer = VideoBuffer::new(10);
let mut buffer = ShadowCaptureVideoBuffer::new(10);

let dummy_frames = [
VideoFrameData::new(vec![1], true, 0),
Expand Down
75 changes: 45 additions & 30 deletions src/encoders/nvenc_encoder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,42 +1,44 @@
use ffmpeg_next::{self as ffmpeg, Rational};
use ringbuf::{
traits::{Producer, Split},
HeapCons, HeapProd, HeapRb,
};

use crate::{application_config::QualityPreset, RawVideoFrame};

use super::{
buffer::{VideoBuffer, VideoFrameData},
video_encoder::{VideoEncoder, GOP_SIZE, ONE_MICROS},
buffer::VideoFrameData,
video_encoder::{VideoEncoder, GOP_SIZE},
};

pub struct NvencEncoder {
encoder: Option<ffmpeg::codec::encoder::Video>,
video_buffer: VideoBuffer,
width: u32,
height: u32,
encoder_name: String,
quality: QualityPreset,
encoded_frame_recv: Option<HeapCons<(i64, VideoFrameData)>>,
encoded_frame_sender: Option<HeapProd<(i64, VideoFrameData)>>,
}

impl VideoEncoder for NvencEncoder {
fn new(
width: u32,
height: u32,
max_buffer_seconds: u32,
quality: QualityPreset,
) -> anyhow::Result<Self>
fn new(width: u32, height: u32, quality: QualityPreset) -> anyhow::Result<Self>
where
Self: Sized,
{
let encoder_name = "h264_nvenc";
let encoder = Some(Self::create_encoder(width, height, encoder_name, &quality)?);
let max_time = max_buffer_seconds as usize * ONE_MICROS;
let video_ring_buffer = HeapRb::<(i64, VideoFrameData)>::new(120);
let (video_ring_sender, video_ring_receiver) = video_ring_buffer.split();

Ok(Self {
encoder,
video_buffer: VideoBuffer::new(max_time),
width,
height,
encoder_name: encoder_name.to_string(),
quality,
encoded_frame_recv: Some(video_ring_receiver),
encoded_frame_sender: Some(video_ring_sender),
})
}

Expand All @@ -56,14 +58,21 @@ impl VideoEncoder for NvencEncoder {
let mut packet = ffmpeg::codec::packet::Packet::empty();
if encoder.receive_packet(&mut packet).is_ok() {
if let Some(data) = packet.data() {
let frame_data = VideoFrameData::new(
data.to_vec(),
packet.is_key(),
packet.pts().unwrap_or(0),
);

self.video_buffer
.insert(packet.dts().unwrap_or(0), frame_data);
if let Some(ref mut sender) = self.encoded_frame_sender {
if sender
.try_push((
packet.dts().unwrap_or(0),
VideoFrameData::new(
data.to_vec(),
packet.is_key(),
packet.pts().unwrap_or(0),
),
))
.is_err()
{
log::error!("Could not send encoded packet to the ringbuf");
}
}
};
}
}
Expand All @@ -77,14 +86,21 @@ impl VideoEncoder for NvencEncoder {
let mut packet = ffmpeg::codec::packet::Packet::empty();
while encoder.receive_packet(&mut packet).is_ok() {
if let Some(data) = packet.data() {
let frame_data = VideoFrameData::new(
data.to_vec(),
packet.is_key(),
packet.pts().unwrap_or(0),
);

self.video_buffer
.insert(packet.dts().unwrap_or(0), frame_data);
if let Some(ref mut sender) = self.encoded_frame_sender {
if sender
.try_push((
packet.dts().unwrap_or(0),
VideoFrameData::new(
data.to_vec(),
packet.is_key(),
packet.pts().unwrap_or(0),
),
))
.is_err()
{
log::error!("Could not send encoded packet to the ringbuf");
}
}
};
packet = ffmpeg::codec::packet::Packet::empty();
}
Expand All @@ -93,7 +109,6 @@ impl VideoEncoder for NvencEncoder {
}

fn drop_encoder(&mut self) {
self.video_buffer.reset();
self.encoder.take();
}

Expand All @@ -112,8 +127,8 @@ impl VideoEncoder for NvencEncoder {
&self.encoder
}

fn get_buffer(&self) -> &VideoBuffer {
&self.video_buffer
fn take_encoded_recv(&mut self) -> Option<HeapCons<(i64, VideoFrameData)>> {
self.encoded_frame_recv.take()
}
}

Expand Down
82 changes: 51 additions & 31 deletions src/encoders/vaapi_encoder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,14 @@ use ffmpeg_next::{
Rational,
};
use log::error;
use ringbuf::{
traits::{Producer, Split},
HeapCons, HeapProd, HeapRb,
};

use super::{
buffer::{VideoBuffer, VideoFrameData},
video_encoder::{VideoEncoder, GOP_SIZE, ONE_MICROS},
buffer::VideoFrameData,
video_encoder::{VideoEncoder, GOP_SIZE},
};

thread_local! {
Expand All @@ -25,32 +29,35 @@ thread_local! {

pub struct VaapiEncoder {
encoder: Option<ffmpeg::codec::encoder::Video>,
video_buffer: VideoBuffer,
width: u32,
height: u32,
encoder_name: String,
quality: QualityPreset,
encoded_frame_recv: Option<HeapCons<(i64, VideoFrameData)>>,
encoded_frame_sender: Option<HeapProd<(i64, VideoFrameData)>>,
}

impl VideoEncoder for VaapiEncoder {
fn new(
width: u32,
height: u32,
max_buffer_seconds: u32,
quality: QualityPreset,
) -> anyhow::Result<Self>
fn new(width: u32, height: u32, quality: QualityPreset) -> anyhow::Result<Self>
where
Self: Sized,
{
// TODO: Should create a child thread that polls the encoder and does the buffer logic
// doing it all at once takes too long

let encoder_name = "h264_vaapi";
let encoder = Self::create_encoder(width, height, encoder_name, &quality)?;
let video_ring_buffer = HeapRb::<(i64, VideoFrameData)>::new(120);
let (video_ring_sender, video_ring_receiver) = video_ring_buffer.split();

Ok(Self {
encoder: Some(encoder),
video_buffer: VideoBuffer::new(max_buffer_seconds as usize * ONE_MICROS),
width,
height,
encoder_name: encoder_name.to_string(),
quality,
encoded_frame_recv: Some(video_ring_receiver),
encoded_frame_sender: Some(video_ring_sender),
})
}

Expand Down Expand Up @@ -126,14 +133,21 @@ impl VideoEncoder for VaapiEncoder {
let mut packet = ffmpeg::codec::packet::Packet::empty();
if encoder.receive_packet(&mut packet).is_ok() {
if let Some(data) = packet.data() {
let frame_data = VideoFrameData::new(
data.to_vec(),
packet.is_key(),
packet.pts().unwrap_or(0),
);

self.video_buffer
.insert(packet.dts().unwrap_or(0), frame_data);
if let Some(ref mut sender) = self.encoded_frame_sender {
if sender
.try_push((
packet.dts().unwrap_or(0),
VideoFrameData::new(
data.to_vec(),
packet.is_key(),
packet.pts().unwrap_or(0),
),
))
.is_err()
{
log::error!("Could not send encoded packet to the ringbuf");
}
}
};
}
});
Expand All @@ -148,14 +162,21 @@ impl VideoEncoder for VaapiEncoder {
let mut packet = ffmpeg::codec::packet::Packet::empty();
while encoder.receive_packet(&mut packet).is_ok() {
if let Some(data) = packet.data() {
let frame_data = VideoFrameData::new(
data.to_vec(),
packet.is_key(),
packet.pts().unwrap_or(0),
);

self.video_buffer
.insert(packet.dts().unwrap_or(0), frame_data);
if let Some(ref mut sender) = self.encoded_frame_sender {
if sender
.try_push((
packet.dts().unwrap_or(0),
VideoFrameData::new(
data.to_vec(),
packet.is_key(),
packet.pts().unwrap_or(0),
),
))
.is_err()
{
log::error!("Could not send encoded packet to the ringbuf");
}
}
};
packet = ffmpeg::codec::packet::Packet::empty();
}
Expand All @@ -178,14 +199,13 @@ impl VideoEncoder for VaapiEncoder {
&self.encoder
}

fn get_buffer(&self) -> &VideoBuffer {
&self.video_buffer
}

fn drop_encoder(&mut self) {
self.video_buffer.reset();
self.encoder.take();
}

fn take_encoded_recv(&mut self) -> Option<HeapCons<(i64, VideoFrameData)>> {
self.encoded_frame_recv.take()
}
}

impl VaapiEncoder {
Expand Down
14 changes: 5 additions & 9 deletions src/encoders/video_encoder.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,26 +1,22 @@
use anyhow::Result;
use ffmpeg_next::{self as ffmpeg};
use ringbuf::HeapCons;

use crate::{application_config::QualityPreset, RawVideoFrame};

use super::buffer::VideoBuffer;
use super::buffer::VideoFrameData;

pub const ONE_MICROS: usize = 1_000_000;
pub const GOP_SIZE: u32 = 30;

pub trait VideoEncoder {
fn new(
width: u32,
height: u32,
max_buffer_seconds: u32,
quality: QualityPreset,
) -> Result<Self>
pub trait VideoEncoder: Send {
fn new(width: u32, height: u32, quality: QualityPreset) -> Result<Self>
where
Self: Sized;
fn process(&mut self, frame: &RawVideoFrame) -> Result<(), ffmpeg::Error>;
fn drain(&mut self) -> Result<(), ffmpeg::Error>;
fn reset(&mut self) -> Result<()>;
fn get_buffer(&self) -> &VideoBuffer;
fn drop_encoder(&mut self);
fn get_encoder(&self) -> &Option<ffmpeg::codec::encoder::Video>;
fn take_encoded_recv(&mut self) -> Option<HeapCons<(i64, VideoFrameData)>>;
}
4 changes: 2 additions & 2 deletions src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ use anyhow::{Context, Error, Result};
use application_config::load_or_create_config;
use encoders::{
audio_encoder::{AudioEncoderImpl, FfmpegAudioEncoder},
buffer::{AudioBuffer, VideoBuffer},
buffer::{AudioBuffer, ShadowCaptureVideoBuffer},
video_encoder::ONE_MICROS,
};
use ffmpeg_next::{self as ffmpeg};
Expand Down Expand Up @@ -82,7 +82,7 @@ async fn main() -> Result<(), Error> {

fn save_buffer(
filename: &str,
video_buffer: &VideoBuffer,
video_buffer: &ShadowCaptureVideoBuffer,
video_encoder: &ffmpeg::codec::encoder::Video,
audio_buffer: &AudioBuffer,
audio_encoder: &FfmpegAudioEncoder,
Expand Down
Loading
Loading