From 7adc1db6720c97be5bec4978d8098ea0b8a09a63 Mon Sep 17 00:00:00 2001 From: enricobuehler Date: Tue, 28 Jul 2026 18:28:46 +0200 Subject: [PATCH] refactor(encode): AVBufferRef ownership moves into an RAII handle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `CudaHw::new` and `VaapiHw::new` are the same shape — alloc a hwdevice, deref it to fill fields, init, alloc a frames ctx, deref, init — and both unwound by hand: `av_buffer_unref` on every failure branch (three in CUDA, two in VAAPI) plus a hand-written `Drop` repeating the pair. That shape has two failure modes and the compiler can see neither: add a branch and forget the cleanup (leak), or let two cleanup paths run (double-unref, an abort inside glibc). `libav::AvBuffer` owns the ref instead. It null-checks on the way in — the check each caller open-coded — and unrefs exactly once on drop, so an early `?` releases whatever was built so far and the failure branches carry no cleanup at all. Both hand-written `Drop` impls are gone, and `linux/mod.rs` goes from seven `av_buffer_unref` calls to zero. The one subtlety, called out at both structs: these two fields must stay declared frames-BEFORE-device. Fields drop in declaration order, and the code being replaced deliberately unref'd frames first (a frames ctx holds its own reference on its device). Refcounting makes either order sound, but a field reorder should not silently change what ships, so the ordering is load-bearing and commented as such. Also adds `cuda_hw_alloc_drop_cycles` — the RAII path had NO test coverage: the NVENC smoke tests take the CPU path and never construct a `CudaHw`, and the VAAPI twin's tests need AMD/Intel silicon. Looping construct/drop is what catches the double-unref (abort) and the leak (allocator growth) this refactor is about. Verified on Nobara (RTX 5070 Ti): `cargo check -p pf-encode --all-targets` clean at exit 0, `cargo test -p pf-encode` 33 passed / 0 failed, and `cuda_hw_alloc_drop_cycles` passes against a real CUDA device — eight construct/drop cycles, no abort. `VaapiHw` is compile-verified only; it is the identical shape but no AMD/Intel box was reachable to run its two ignored tests. --- crates/pf-encode/src/enc/libav.rs | 46 ++++++++++++ crates/pf-encode/src/enc/linux/mod.rs | 95 +++++++++++++++---------- crates/pf-encode/src/enc/linux/vaapi.rs | 60 +++++++--------- 3 files changed, 130 insertions(+), 71 deletions(-) diff --git a/crates/pf-encode/src/enc/libav.rs b/crates/pf-encode/src/enc/libav.rs index 71cbd6e2..94710e69 100644 --- a/crates/pf-encode/src/enc/libav.rs +++ b/crates/pf-encode/src/enc/libav.rs @@ -26,6 +26,52 @@ pub(crate) fn pixel_to_av(p: Pixel) -> ffi::AVPixelFormat { ffi::AVPixelFormat::from(p) } +/// An owned `AVBufferRef` — unref'd exactly once, when it drops. +/// +/// The hwdevice/hwframes constructors used to unref by hand on *every* failure branch (three in the +/// CUDA path, two in VAAPI) and then once more in a hand-written `Drop`. That shape leaks the moment +/// somebody adds a branch and forgets the cleanup, and double-unrefs the moment two of them run — +/// and neither mistake is visible at the call site or catchable by the compiler. Ownership lives in +/// this type instead: an early `?` drops whatever was built so far, in reverse construction order, +/// with no cleanup code at the call site at all. +/// +/// **Drop order** matters to the callers holding two of these. A frames context internally holds its +/// own reference on its device, and the code this replaced deliberately unref'd frames *before* +/// device. Rust drops struct fields in DECLARATION order, so a struct holding both must declare +/// frames before device to keep that. Refcounting makes either order sound in principle — +/// the device cannot die while a frames ctx still references it — but the observable order is kept +/// exactly as it shipped rather than quietly inverted by a field reorder. +pub(crate) struct AvBuffer(*mut ffi::AVBufferRef); + +impl AvBuffer { + /// Take ownership of a freshly-created `AVBufferRef`, rejecting the null that an ffmpeg + /// allocator returns on failure (so the `is_null` check every caller used to open-code happens + /// once, here). + /// + /// # Safety + /// `p` must be null, or a live `AVBufferRef` whose ownership passes to the returned value — + /// nothing else may unref it. + pub(crate) unsafe fn from_raw(p: *mut ffi::AVBufferRef) -> Option { + (!p.is_null()).then_some(AvBuffer(p)) + } + + /// The borrowed pointer, for the ffmpeg calls that read a ref without consuming it. Borrowed + /// only — the `AvBuffer` stays the owner, so callers must not unref what this returns. + pub(crate) fn as_ptr(&self) -> *mut ffi::AVBufferRef { + self.0 + } +} + +impl Drop for AvBuffer { + fn drop(&mut self) { + // SAFETY: `self.0` is the non-null ref `from_raw` took ownership of, and this type is its + // sole owner (it is neither `Clone` nor `Copy`, and `as_ptr` only lends), so this runs + // exactly once for that reference. `av_buffer_unref` drops the one reference and nulls the + // pointer through the `&mut`. + unsafe { ffi::av_buffer_unref(&mut self.0) }; + } +} + /// One `receive_packet` attempt, with the not-ready states kept distinct so a blocking drain can /// tell "still encoding" (retry) from "stream over" (stop). The Linux NVENC/VAAPI polls collapse /// `Again`/`Eof` to `None`; the Windows AMF/QSV path keeps them apart for its deadline-driven loop. diff --git a/crates/pf-encode/src/enc/linux/mod.rs b/crates/pf-encode/src/enc/linux/mod.rs index a223f178..50d30857 100644 --- a/crates/pf-encode/src/enc/linux/mod.rs +++ b/crates/pf-encode/src/enc/linux/mod.rs @@ -22,7 +22,8 @@ use std::os::raw::c_int; use std::ptr; use super::libav::{ - apply_low_latency_rc, pixel_to_av, poll_encoder, PollOutcome, SWS_CS_ITU709, SWS_POINT, + apply_low_latency_rc, pixel_to_av, poll_encoder, AvBuffer, PollOutcome, SWS_CS_ITU709, + SWS_POINT, }; use ffmpeg::ffi; // = ffmpeg_sys_next @@ -60,8 +61,11 @@ struct AVCUDADeviceContext { /// CUDA hardware-frame contexts that wrap our shared `CUcontext`, so `hevc_nvenc` reads the /// imported device buffer directly. Owns two `AVBufferRef`s, unref'd on drop. struct CudaHw { - device_ref: *mut ffi::AVBufferRef, - frames_ref: *mut ffi::AVBufferRef, + // Declared frames-BEFORE-device on purpose: these drop in declaration order, and that + // reproduces exactly what the hand-written `Drop` this replaced did (the frames ctx holds its + // own reference on the device). Do not reorder these two fields. + frames_ref: AvBuffer, + device_ref: AvBuffer, } impl CudaHw { @@ -72,57 +76,42 @@ impl CudaHw { /// (`nvenc_open_einval`), and a hwdevice/hwframes EINVAL is a config error no bitrate can /// fix — enrolling it would burn ~10 doomed encoder opens before surfacing the real failure. unsafe fn new(cu_ctx: *mut std::ffi::c_void, sw_format: Pixel, w: u32, h: u32) -> Result { - let mut device_ref = ffi::av_hwdevice_ctx_alloc(ffi::AVHWDeviceType::AV_HWDEVICE_TYPE_CUDA); - if device_ref.is_null() { - bail!("av_hwdevice_ctx_alloc(CUDA) failed"); - } - let dev_ctx = (*device_ref).data as *mut ffi::AVHWDeviceContext; + // Each `?`/`bail!` below drops whatever has been built so far — `AvBuffer`'s `Drop` is the + // single unref path, so the failure branches carry no cleanup of their own. + let device_ref = AvBuffer::from_raw(ffi::av_hwdevice_ctx_alloc( + ffi::AVHWDeviceType::AV_HWDEVICE_TYPE_CUDA, + )) + .context("av_hwdevice_ctx_alloc(CUDA) failed")?; + let dev_ctx = (*device_ref.as_ptr()).data as *mut ffi::AVHWDeviceContext; let cu = (*dev_ctx).hwctx as *mut AVCUDADeviceContext; (*cu).cuda_ctx = cu_ctx; // share the importer's context - let r = ffi::av_hwdevice_ctx_init(device_ref); + let r = ffi::av_hwdevice_ctx_init(device_ref.as_ptr()); if r < 0 { - ffi::av_buffer_unref(&mut device_ref); bail!("av_hwdevice_ctx_init failed ({r})"); } - let mut frames_ref = ffi::av_hwframe_ctx_alloc(device_ref); - if frames_ref.is_null() { - ffi::av_buffer_unref(&mut device_ref); - bail!("av_hwframe_ctx_alloc failed"); - } - let fc = (*frames_ref).data as *mut ffi::AVHWFramesContext; + let frames_ref = AvBuffer::from_raw(ffi::av_hwframe_ctx_alloc(device_ref.as_ptr())) + .context("av_hwframe_ctx_alloc failed")?; + let fc = (*frames_ref.as_ptr()).data as *mut ffi::AVHWFramesContext; (*fc).format = ffi::AVPixelFormat::AV_PIX_FMT_CUDA; (*fc).sw_format = pixel_to_av(sw_format); (*fc).width = w as c_int; (*fc).height = h as c_int; (*fc).initial_pool_size = 0; // we supply the device pointers - let r = ffi::av_hwframe_ctx_init(frames_ref); + let r = ffi::av_hwframe_ctx_init(frames_ref.as_ptr()); if r < 0 { - ffi::av_buffer_unref(&mut frames_ref); - ffi::av_buffer_unref(&mut device_ref); bail!("av_hwframe_ctx_init failed ({r})"); } Ok(CudaHw { - device_ref, frames_ref, + device_ref, }) } } -impl Drop for CudaHw { - fn drop(&mut self) { - // SAFETY: `frames_ref`/`device_ref` are the two non-null `AVBufferRef`s `CudaHw::new` created - // (it bails before returning `Self` if either alloc fails, so a live `CudaHw` always holds - // both). `av_buffer_unref` drops one reference and nulls the pointer through the `&mut`. This - // `Drop` runs exactly once and `CudaHw` owns these refs exclusively → no double-free / - // use-after-free. Frames are unref'd before the device (the frames ctx internally refs the - // device; refcounted, so the order is sound regardless). - unsafe { - ffi::av_buffer_unref(&mut self.frames_ref); - ffi::av_buffer_unref(&mut self.device_ref); - } - } -} +// No `Drop` for `CudaHw`: each `AvBuffer` field unrefs itself, in declaration order (frames, then +// device — see the field comment). The hand-written unref pair this replaced had to be kept in sync +// with every failure branch in `new`; now there is exactly one unref path and it cannot be skipped. /// Map a captured layout to the NVENC input pixel format, and whether a 3→4 byte expand is /// needed (packed RGB/BGR have no padding byte; the NVENC `*0` formats do). @@ -421,8 +410,8 @@ impl NvencEncoder { unsafe { let raw = video.as_mut_ptr(); (*raw).pix_fmt = ffi::AVPixelFormat::AV_PIX_FMT_CUDA; - (*raw).hw_device_ctx = ffi::av_buffer_ref(hw.device_ref); - (*raw).hw_frames_ctx = ffi::av_buffer_ref(hw.frames_ref); + (*raw).hw_device_ctx = ffi::av_buffer_ref(hw.device_ref.as_ptr()); + (*raw).hw_frames_ctx = ffi::av_buffer_ref(hw.frames_ref.as_ptr()); } Some(hw) } else { @@ -847,7 +836,8 @@ impl NvencEncoder { .cuda .as_ref() .context("CUDA hw context missing (encoder opened in CPU mode)")? - .frames_ref; + .frames_ref + .as_ptr(); // The device→device copy below uses our shared context directly; make it current on the // encode thread (ffmpeg pushes its own around the pool alloc, so order is fine). pf_zerocopy::cuda::make_current().context("CUDA context current (encode thread)")?; @@ -1055,6 +1045,37 @@ pub fn probe_can_encode_10bit(codec: Codec) -> bool { .is_ok() } +#[cfg(test)] +mod cuda_hw_tests { + use super::*; + + /// `CudaHw` owns its two `AVBufferRef`s through `AvBuffer`, so *construct and drop* is the + /// entire contract: a missed unref leaks, a doubled one aborts inside glibc. Nothing else in + /// the suite covers it — the NVENC smoke tests take the CPU path and never build one, and the + /// VAAPI twin's tests need AMD/Intel silicon. Looping the cycle is the point: a double-unref + /// shows up as an abort, and a leak shows as the allocator growing across iterations. + /// + /// `#[ignore]`d (needs a real CUDA device): + /// `cargo test -p pf-encode cuda_hw_alloc_drop_cycles -- --ignored --nocapture` + #[test] + #[ignore = "needs a real CUDA device (run on an NVIDIA host, not the build box)"] + fn cuda_hw_alloc_drop_cycles() { + ffmpeg::init().expect("libav init"); + let cu_ctx = pf_zerocopy::cuda::context().expect("shared CUDA context"); + for i in 0..8 { + // SAFETY: `CudaHw::new` requires libav initialized (asserted above) and a valid + // `CUcontext` — `cu_ctx` is the live shared context from `pf_zerocopy`. NV12 at + // 640x480 are a valid format and positive dims. The handle drops at the end of each + // iteration, which is precisely the unref path under test. + let hw = unsafe { CudaHw::new(cu_ctx.cast(), Pixel::NV12, 640, 480) } + .unwrap_or_else(|e| panic!("CudaHw::new failed on iteration {i}: {e:#}")); + assert!(!hw.device_ref.as_ptr().is_null(), "device ref went null"); + assert!(!hw.frames_ref.as_ptr().is_null(), "frames ref went null"); + } + eprintln!("8 CudaHw alloc/drop cycles completed without abort"); + } +} + #[cfg(test)] mod hdr_tests { use super::*; diff --git a/crates/pf-encode/src/enc/linux/vaapi.rs b/crates/pf-encode/src/enc/linux/vaapi.rs index fc9e2f44..5449030e 100644 --- a/crates/pf-encode/src/enc/linux/vaapi.rs +++ b/crates/pf-encode/src/enc/linux/vaapi.rs @@ -36,7 +36,8 @@ use std::ptr; use std::sync::{Mutex, OnceLock}; use super::libav::{ - apply_low_latency_rc, pixel_to_av, poll_encoder, PollOutcome, SWS_CS_ITU709, SWS_POINT, + apply_low_latency_rc, pixel_to_av, poll_encoder, AvBuffer, PollOutcome, SWS_CS_ITU709, + SWS_POINT, }; use ffmpeg::ffi; // = ffmpeg_sys_next @@ -396,8 +397,8 @@ pub fn probe_can_encode(codec: Codec) -> bool { 480, 30, 2_000_000, - hw.device_ref, - hw.frames_ref, + hw.device_ref.as_ptr(), + hw.frames_ref.as_ptr(), false, ) .is_ok(), @@ -435,8 +436,8 @@ pub fn probe_can_encode_10bit(codec: Codec) -> bool { 480, 30, 2_000_000, - hw.device_ref, - hw.frames_ref, + hw.device_ref.as_ptr(), + hw.frames_ref.as_ptr(), true, ) .is_ok(), @@ -463,8 +464,11 @@ pub fn probe_can_encode_444(_codec: Codec) -> bool { /// VAAPI device + NV12 frames pool (the encoder's input surfaces for the CPU path). struct VaapiHw { - device_ref: *mut ffi::AVBufferRef, - frames_ref: *mut ffi::AVBufferRef, + // Declared frames-BEFORE-device on purpose: these drop in declaration order, which reproduces + // exactly what the hand-written `Drop` this replaced did (the frames ctx holds its own + // reference on the device). Do not reorder these two fields. + frames_ref: AvBuffer, + device_ref: AvBuffer, } impl VaapiHw { @@ -481,44 +485,32 @@ impl VaapiHw { if r < 0 { bail!("no VAAPI device ({:?}): {}", node, ffmpeg::Error::from(r)); } - let mut frames_ref = ffi::av_hwframe_ctx_alloc(device_ref); - if frames_ref.is_null() { - ffi::av_buffer_unref(&mut device_ref); - bail!("av_hwframe_ctx_alloc(VAAPI) failed"); - } - let fc = (*frames_ref).data as *mut ffi::AVHWFramesContext; + // `av_hwdevice_ctx_create` wrote an owned ref into `device_ref`; take ownership of it here so + // every `bail!` below drops it (and the frames ref, once built) without cleanup of its own. + let device_ref = AvBuffer::from_raw(device_ref) + .context("av_hwdevice_ctx_create(VAAPI) gave no device")?; + let frames_ref = AvBuffer::from_raw(ffi::av_hwframe_ctx_alloc(device_ref.as_ptr())) + .context("av_hwframe_ctx_alloc(VAAPI) failed")?; + let fc = (*frames_ref.as_ptr()).data as *mut ffi::AVHWFramesContext; (*fc).format = ffi::AVPixelFormat::AV_PIX_FMT_VAAPI; (*fc).sw_format = sw_format; (*fc).width = w as c_int; (*fc).height = h as c_int; (*fc).initial_pool_size = pool; - let r = ffi::av_hwframe_ctx_init(frames_ref); + let r = ffi::av_hwframe_ctx_init(frames_ref.as_ptr()); if r < 0 { - ffi::av_buffer_unref(&mut frames_ref); - ffi::av_buffer_unref(&mut device_ref); bail!("av_hwframe_ctx_init(VAAPI) failed ({r})"); } Ok(VaapiHw { - device_ref, frames_ref, + device_ref, }) } } -impl Drop for VaapiHw { - fn drop(&mut self) { - // SAFETY: `frames_ref`/`device_ref` are the two non-null `AVBufferRef`s `VaapiHw::new` - // created (it bails before constructing `Self` if either alloc fails, so a live `VaapiHw` - // always holds both). `av_buffer_unref` drops one reference and nulls the pointer through the - // `&mut`. This `Drop` runs exactly once and `VaapiHw` owns these refs exclusively, so there - // is no double-free / use-after-free. Frames are unref'd before the device because the frames - // ctx internally holds a ref on the device (refcounted, so the order is sound either way). - unsafe { - ffi::av_buffer_unref(&mut self.frames_ref); - ffi::av_buffer_unref(&mut self.device_ref); - } - } -} +// No `Drop` for `VaapiHw`: each `AvBuffer` field unrefs itself, in declaration order (frames, then +// device — see the field comment). The hand-written unref pair this replaced had to stay in sync +// with every failure branch in `new`; now there is one unref path and it cannot be skipped. struct CpuInner { enc: encoder::video::Encoder, @@ -567,8 +559,8 @@ impl CpuInner { height, fps, bitrate_bps, - hw.device_ref, - hw.frames_ref, + hw.device_ref.as_ptr(), + hw.frames_ref.as_ptr(), ten_bit, )? }; @@ -698,7 +690,7 @@ impl CpuInner { if hwf.is_null() { bail!("av_frame_alloc(hw) failed"); } - if ffi::av_hwframe_get_buffer(self.hw.frames_ref, hwf, 0) < 0 { + if ffi::av_hwframe_get_buffer(self.hw.frames_ref.as_ptr(), hwf, 0) < 0 { ffi::av_frame_free(&mut hwf); bail!("av_hwframe_get_buffer(VAAPI) failed"); }