use std::ffi::CString;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use crossbeam_channel::{SendTimeoutError, Sender};
use ffmpeg_sys_next::AVPixelFormat::AV_PIX_FMT_NONE;
use ffmpeg_sys_next::{
av_get_pix_fmt, av_image_get_buffer_size, av_pix_fmt_desc_get, AVPixelFormat,
AV_PIX_FMT_FLAG_HWACCEL,
};
use crate::core::context::ffmpeg_context::build_writer_context;
use crate::core::context::frame_source::FrameSourceParams;
use crate::core::context::output::Output;
use crate::core::scheduler::ffmpeg_scheduler::{is_stopping, FfmpegScheduler, Running};
use crate::error::OpenOutputError;
const QUEUE_BUDGET_BYTES: usize = 64 * 1024 * 1024;
const SEND_POLL: Duration = Duration::from_millis(100);
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum WriterError {
#[error("invalid dimensions {width}x{height}")]
InvalidDimensions { width: u32, height: u32 },
#[error("unknown pixel format '{0}'")]
UnknownPixelFormat(String),
#[error("hardware pixel format '{0}' cannot be pushed from CPU memory")]
HardwarePixelFormat(String),
#[error("invalid fps {num}/{den}: both must be positive")]
InvalidFps { num: i32, den: i32 },
#[error("queue_capacity must be >= 1")]
ZeroQueueCapacity,
#[error("output consumes no video stream from the pushed frames")]
NoVideoDestination,
#[error(
"filter_desc must have exactly one video input pad and one video \
output pad; found {input_pads} input pad(s) ({video_input_pads} \
video) and {output_pads} output pad(s) ({video_output_pads} video)"
)]
FilterShape {
input_pads: usize,
video_input_pads: usize,
output_pads: usize,
video_output_pads: usize,
},
#[error("filter_desc must be a single connected graph; found {components} disconnected parts")]
DisconnectedFilterGraph { components: usize },
#[error(
"filter_desc has no directed path from its input pad to its output \
pad; the pushed frames could not influence the encoded output"
)]
UnreachableFilterOutput,
#[error("stream maps are not supported by VideoWriter; the pushed frames are the only stream")]
StreamMapsUnsupported,
}
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum PushError {
#[error("frame has {got} bytes, expected exactly {expected} (tightly packed)")]
InvalidSize { expected: usize, got: usize },
#[error("pipeline closed; call finish() to retrieve the pipeline result")]
PipelineClosed,
}
pub struct VideoWriter {
sender: Option<Sender<Vec<u8>>>,
scheduler: Option<FfmpegScheduler<Running>>,
frame_size: usize,
status: Arc<AtomicUsize>,
}
impl VideoWriter {
pub fn builder(width: u32, height: u32) -> VideoWriterBuilder {
VideoWriterBuilder::new(width, height)
}
pub fn frame_size(&self) -> usize {
self.frame_size
}
pub fn write(&mut self, frame: &[u8]) -> Result<(), PushError> {
if frame.len() != self.frame_size {
return Err(PushError::InvalidSize {
expected: self.frame_size,
got: frame.len(),
});
}
self.enqueue(frame.to_vec())
}
pub fn write_owned(&mut self, frame: Vec<u8>) -> Result<(), PushError> {
if frame.len() != self.frame_size {
return Err(PushError::InvalidSize {
expected: self.frame_size,
got: frame.len(),
});
}
self.enqueue(frame)
}
fn enqueue(&mut self, frame: Vec<u8>) -> Result<(), PushError> {
if is_stopping(self.status.load(Ordering::Acquire)) {
return Err(PushError::PipelineClosed);
}
let sender = match &self.sender {
Some(sender) => sender,
None => return Err(PushError::PipelineClosed),
};
let mut msg = frame;
loop {
match sender.send_timeout(msg, SEND_POLL) {
Ok(()) => return Ok(()),
Err(SendTimeoutError::Timeout(returned)) => {
if is_stopping(self.status.load(Ordering::Acquire)) {
return Err(PushError::PipelineClosed);
}
msg = returned;
}
Err(SendTimeoutError::Disconnected(_)) => return Err(PushError::PipelineClosed),
}
}
}
pub fn finish(mut self) -> crate::error::Result<()> {
self.sender = None; match self.scheduler.take() {
Some(scheduler) => scheduler.wait(),
None => Ok(()),
}
}
pub fn abort(mut self) {
self.sender = None;
if let Some(scheduler) = self.scheduler.take() {
scheduler.abort();
}
}
}
impl Drop for VideoWriter {
fn drop(&mut self) {
self.sender = None;
if let Some(scheduler) = self.scheduler.take() {
if let Err(e) = scheduler.wait() {
log::error!("VideoWriter dropped without finish(); pipeline error: {e}");
}
}
}
}
pub struct VideoWriterBuilder {
width: u32,
height: u32,
pixel_format: String,
fps_num: i32,
fps_den: i32,
queue_capacity: Option<usize>,
filter_desc: Option<String>,
}
impl VideoWriterBuilder {
fn new(width: u32, height: u32) -> Self {
Self {
width,
height,
pixel_format: "rgba".to_string(),
fps_num: 30,
fps_den: 1,
queue_capacity: None,
filter_desc: None,
}
}
pub fn pixel_format(mut self, fmt: impl Into<String>) -> Self {
self.pixel_format = fmt.into();
self
}
pub fn fps(mut self, num: i32, den: i32) -> Self {
self.fps_num = num;
self.fps_den = den;
self
}
pub fn queue_capacity(mut self, frames: usize) -> Self {
self.queue_capacity = Some(frames);
self
}
pub fn filter_desc(mut self, desc: impl Into<String>) -> Self {
self.filter_desc = Some(desc.into());
self
}
pub fn open(self, output: impl Into<Output>) -> crate::error::Result<VideoWriter> {
if self.width == 0
|| self.height == 0
|| self.width > i32::MAX as u32
|| self.height > i32::MAX as u32
{
return Err(WriterError::InvalidDimensions {
width: self.width,
height: self.height,
}
.into());
}
if self.fps_num <= 0 || self.fps_den <= 0 {
return Err(WriterError::InvalidFps {
num: self.fps_num,
den: self.fps_den,
}
.into());
}
let (pix_fmt, frame_size) =
resolve_source_format(&self.pixel_format, self.width, self.height)?;
let queue_capacity = match self.queue_capacity {
Some(0) => return Err(WriterError::ZeroQueueCapacity.into()),
Some(n) => n,
None => adaptive_capacity(frame_size),
};
let params = FrameSourceParams {
width: self.width as i32,
height: self.height as i32,
pix_fmt,
fps_num: self.fps_num,
fps_den: self.fps_den,
};
let (ctx, sender) = match build_writer_context(
params,
queue_capacity,
self.filter_desc.as_deref(),
output.into(),
) {
Ok(built) => built,
Err(crate::error::Error::OpenOutput(OpenOutputError::NotContainStream)) => {
return Err(WriterError::NoVideoDestination.into());
}
Err(e) => return Err(e),
};
let status = ctx.scheduler_status.clone();
let scheduler = FfmpegScheduler::new(ctx);
let scheduler = scheduler.start()?;
Ok(VideoWriter {
sender: Some(sender),
scheduler: Some(scheduler),
frame_size,
status,
})
}
}
fn adaptive_capacity(frame_size: usize) -> usize {
(QUEUE_BUDGET_BYTES / frame_size).clamp(1, 4)
}
fn resolve_source_format(
pixel_format: &str,
width: u32,
height: u32,
) -> Result<(AVPixelFormat, usize), WriterError> {
let cstr = CString::new(pixel_format)
.map_err(|_| WriterError::UnknownPixelFormat(pixel_format.to_string()))?;
let pix_fmt = unsafe { av_get_pix_fmt(cstr.as_ptr()) };
if pix_fmt == AV_PIX_FMT_NONE {
return Err(WriterError::UnknownPixelFormat(pixel_format.to_string()));
}
let desc = unsafe { av_pix_fmt_desc_get(pix_fmt) };
if !desc.is_null() && unsafe { (*desc).flags } & (AV_PIX_FMT_FLAG_HWACCEL as u64) != 0 {
return Err(WriterError::HardwarePixelFormat(pixel_format.to_string()));
}
let size = unsafe { av_image_get_buffer_size(pix_fmt, width as i32, height as i32, 1) };
if size <= 0 {
return Err(WriterError::UnknownPixelFormat(pixel_format.to_string()));
}
Ok((pix_fmt, size as usize))
}
#[cfg(test)]
mod tests {
use super::*;
fn frame_size_of(fmt: &str, w: u32, h: u32) -> Result<usize, WriterError> {
resolve_source_format(fmt, w, h).map(|(_, size)| size)
}
#[test]
fn frame_size_matches_ffmpeg_for_common_formats() {
assert_eq!(frame_size_of("rgba", 64, 48).unwrap(), 64 * 48 * 4);
assert_eq!(frame_size_of("rgb24", 64, 48).unwrap(), 64 * 48 * 3);
assert_eq!(frame_size_of("gray8", 64, 48).unwrap(), 64 * 48);
assert_eq!(frame_size_of("yuv420p", 64, 48).unwrap(), 64 * 48 * 3 / 2);
assert_eq!(frame_size_of("nv12", 64, 48).unwrap(), 64 * 48 * 3 / 2);
}
#[test]
fn yuv420p_odd_height_rounds_chroma_up() {
let expected = 64 * 49 + 2 * (32 * 25);
assert_eq!(frame_size_of("yuv420p", 64, 49).unwrap(), expected);
}
#[test]
fn unknown_pixel_format_is_rejected() {
assert!(matches!(
frame_size_of("definitely_not_a_format", 16, 16),
Err(WriterError::UnknownPixelFormat(_))
));
}
#[test]
fn hardware_pixel_format_is_rejected() {
assert!(matches!(
frame_size_of("cuda", 16, 16),
Err(WriterError::HardwarePixelFormat(_)) | Err(WriterError::UnknownPixelFormat(_))
));
}
#[test]
fn adaptive_capacity_scales_with_frame_size() {
assert_eq!(adaptive_capacity(1), 4);
assert_eq!(adaptive_capacity(1024), 4);
assert_eq!(adaptive_capacity(1920 * 1080 * 4), 4);
assert_eq!(adaptive_capacity(3840 * 2160 * 4), 2);
assert_eq!(adaptive_capacity(QUEUE_BUDGET_BYTES * 2), 1);
}
#[test]
fn invalid_dimensions_and_fps_are_rejected_at_open() {
use crate::Output;
let out = || Output::from("/dev/null").set_video_codec("mpeg4");
assert!(matches!(
VideoWriter::builder(0, 48).open(out()),
Err(crate::error::Error::Writer(
WriterError::InvalidDimensions { .. }
))
));
assert!(matches!(
VideoWriter::builder(64, 48).fps(0, 1).open(out()),
Err(crate::error::Error::Writer(WriterError::InvalidFps { .. }))
));
assert!(matches!(
VideoWriter::builder(64, 48).queue_capacity(0).open(out()),
Err(crate::error::Error::Writer(WriterError::ZeroQueueCapacity))
));
assert!(matches!(
VideoWriter::builder(64, 48)
.pixel_format("nope")
.open(out()),
Err(crate::error::Error::Writer(
WriterError::UnknownPixelFormat(_)
))
));
}
#[allow(dead_code)]
fn error_composition(writer: &mut VideoWriter, frame: &[u8]) -> crate::error::Result<()> {
writer.write(frame)?;
Ok(())
}
#[test]
fn push_error_display_reports_sizes() {
let e = PushError::InvalidSize {
expected: 12288,
got: 100,
};
assert!(e.to_string().contains("12288"));
assert!(e.to_string().contains("100"));
}
#[test]
fn scheduler_teardown_completes_with_live_ingress_sender() {
use crate::core::context::ffmpeg_context::build_writer_context;
use ffmpeg_sys_next::AVPixelFormat::AV_PIX_FMT_RGBA;
for (label, abort) in [("drop", false), ("abort", true)] {
let out = std::env::temp_dir().join(format!(
"ez_ffmpeg_writer_sender_held_{label}_{}.mp4",
std::process::id()
));
let params = FrameSourceParams {
width: 64,
height: 48,
pix_fmt: AV_PIX_FMT_RGBA,
fps_num: 30,
fps_den: 1,
};
let (ctx, sender) = build_writer_context(
params,
2,
None,
Output::from(out.to_str().unwrap()).set_video_codec("mpeg4"),
)
.expect("writer context");
let scheduler = FfmpegScheduler::new(ctx).start().expect("start");
sender
.send(vec![0u8; 64 * 48 * 4])
.expect("worker must be consuming");
let (tx, rx) = std::sync::mpsc::channel();
std::thread::spawn(move || {
if abort {
scheduler.abort();
} else {
drop(scheduler);
}
let _ = tx.send(());
});
rx.recv_timeout(Duration::from_secs(30))
.unwrap_or_else(|_| panic!("{label}: teardown hung with a live ingress sender"));
drop(sender);
}
}
}