Files
SR-Label-SamTool/sam-tool-tauri/src-tauri/src/decoder.rs
2026-06-04 19:50:02 +05:30

972 lines
34 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
use crate::video::VideoInfo;
use std::collections::{HashMap, VecDeque};
use std::io::{BufRead, BufReader};
use std::path::PathBuf;
use std::process::{Child, ChildStderr, ChildStdout, Command, Stdio};
use std::sync::{Arc, Condvar, Mutex};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};
/// Cap on stderr lines retained per ffmpeg process. The hwaccel-init message
/// is usually 1-3 lines; we keep more to cover later-frame errors too.
const STDERR_TAIL_LINES: usize = 32;
pub const DEFAULT_CACHE_CAPACITY: usize = 160;
pub const MIN_CACHE_CAPACITY: usize = 50;
pub const MAX_CACHE_CAPACITY: usize = 800;
/// Cache cap for the RGBA preview decoder. Raw frames are 30-80× bigger
/// than MJPEG so we cap depth for memory headroom. Sized for 1080p (the
/// most common case): 16 frames × 8.3 MB ≈ 133 MB. At 720p this drops to
/// ~45 MB. At 4K it's ~528 MB, so users with 4K sources should set
/// preview_max_height in settings.
pub const RGBA_PREVIEW_CACHE_CAPACITY: usize = 16;
pub const RGBA_PREVIEW_LOOKAHEAD: u64 = 8;
const READ_BUF_BYTES: usize = 256 * 1024;
/// Wire format produced by an ffmpeg child. Preview decoder uses Rgba so
/// the JS side avoids a CPU MJPEG decode per frame (which was the real
/// per-frame bottleneck — see review findings A + E). Extract decoder
/// stays on Mjpeg because `extract_one` writes the decoder output to disk
/// and users expect .jpg files.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OutputFormat {
Mjpeg,
Rgba,
}
/// If the requested frame is this far ahead of the worker, kill + reseek.
/// Higher = fewer ffmpeg restarts during forward scan; the cap is bounded
/// by lookahead × cache capacity anyway.
const SEEK_AHEAD_THRESHOLD: u64 = 120;
const DECODE_TIMEOUT: Duration = Duration::from_secs(15);
struct FrameCache {
map: HashMap<u64, Vec<u8>>,
order: VecDeque<u64>,
capacity: usize,
}
impl FrameCache {
fn new(capacity: usize) -> Self {
Self {
map: HashMap::new(),
order: VecDeque::new(),
capacity,
}
}
fn contains(&self, idx: u64) -> bool {
self.map.contains_key(&idx)
}
/// FIFO retrieval (no reorder-on-access). Our real access pattern is
/// sequential decode, so cache fills in order and evicts in order; the
/// previous LRU reshuffle was an O(n) VecDeque scan that bought nothing.
fn get(&self, idx: u64) -> Option<Vec<u8>> {
self.map.get(&idx).cloned()
}
fn insert(&mut self, idx: u64, bytes: Vec<u8>) {
if self.map.contains_key(&idx) {
// Re-decode of an already-cached frame — just overwrite, no
// order change (the queue entry stays valid).
self.map.insert(idx, bytes);
return;
}
if self.map.len() >= self.capacity {
if let Some(oldest) = self.order.pop_front() {
self.map.remove(&oldest);
}
}
self.map.insert(idx, bytes);
self.order.push_back(idx);
}
fn set_capacity(&mut self, new_cap: usize) {
self.capacity = new_cap.clamp(MIN_CACHE_CAPACITY, MAX_CACHE_CAPACITY);
while self.map.len() > self.capacity {
if let Some(oldest) = self.order.pop_front() {
self.map.remove(&oldest);
} else {
break;
}
}
}
}
struct SharedState {
/// Owned by shared state (not the worker closure) so stop_worker can
/// kill ffmpeg directly and unblock a pipe read that's waiting for the
/// first post-seek frame.
child: Option<Child>,
cache: FrameCache,
/// Frame the worker's pipe will emit next.
next_idx: u64,
/// Most recent frame the frontend requested — used to bound lookahead.
reader_idx: u64,
/// Worker reads this many frames past reader_idx before parking.
lookahead: u64,
eof: bool,
error: Option<String>,
shutdown: bool,
/// Last N stderr lines from the most recent ffmpeg child. Populated by
/// a non-blocking drain thread; read only when something has gone wrong
/// (e.g. hwaccel fallback) to surface the ffmpeg diagnostic.
stderr_tail: VecDeque<String>,
}
struct Shared {
state: Mutex<SharedState>,
cond: Condvar,
}
pub struct DecoderSession {
path: PathBuf,
info: VideoInfo,
/// ffmpeg `-hwaccel` argument. None / empty = software decode. Cleared
/// on any decode failure so the next attempt uses software — we'd
/// rather degrade than keep retrying a flaky GPU path.
hwaccel: Option<String>,
/// When Some(H), ffmpeg downscales to height H (aspect preserved) via
/// `-vf scale=-2:H`. Used by the playback decoder for cheap previews;
/// the extract decoder always passes None so saved frames are full-res.
format: OutputFormat,
/// Pre-computed (width, height) ffmpeg will emit. For Mjpeg this is the
/// source dims (mjpeg encoder doesn't care). For Rgba these MUST be
/// even and exact, since the reader frames input by `out_w * out_h * 4`
/// byte counts — there's no marker to find frame boundaries in raw.
output_dims: (u32, u32),
shared: Arc<Shared>,
worker: Option<JoinHandle<()>>,
/// One-shot notice surfaced via `take_notice` after a hwaccel fallback.
/// Lets the frontend show "Hardware decode 'cuda' failed; running
/// software" instead of silently testing the same code path under three
/// different settings.
last_notice: Mutex<Option<String>>,
}
impl DecoderSession {
pub fn open(
path: PathBuf,
info: VideoInfo,
hwaccel: Option<String>,
preview_max_height: Option<u32>,
format: OutputFormat,
) -> Self {
let preview_max_height = preview_max_height.filter(|h| *h > 0);
let output_dims = compute_output_dims(&info, preview_max_height);
// RGBA frames are huge (1080p = 8.3 MB), so a 160-entry cache would
// be ~1.3 GB. Cap depth and lookahead for the Rgba path; Mjpeg
// keeps its existing generous defaults.
let (init_cap, init_lookahead) = match format {
OutputFormat::Mjpeg => (
DEFAULT_CACHE_CAPACITY,
(DEFAULT_CACHE_CAPACITY as u64) / 2,
),
OutputFormat::Rgba => (RGBA_PREVIEW_CACHE_CAPACITY, RGBA_PREVIEW_LOOKAHEAD),
};
Self {
path,
info,
hwaccel: hwaccel.filter(|s| !s.is_empty()),
format,
output_dims,
shared: Arc::new(Shared {
state: Mutex::new(SharedState {
child: None,
cache: FrameCache::new(init_cap),
next_idx: 0,
reader_idx: 0,
lookahead: init_lookahead,
eof: false,
error: None,
shutdown: false,
stderr_tail: VecDeque::new(),
}),
cond: Condvar::new(),
}),
worker: None,
last_notice: Mutex::new(None),
}
}
/// Drain and return the most recent fallback / diagnostic notice (if any).
/// One-shot: a subsequent call returns None until the next event sets it.
pub fn take_notice(&self) -> Option<String> {
self.last_notice.lock().ok().and_then(|mut n| n.take())
}
fn snapshot_stderr_tail(&self) -> String {
let st = self.shared.state.lock().unwrap();
if st.stderr_tail.is_empty() {
"(no ffmpeg stderr captured)".into()
} else {
st.stderr_tail
.iter()
.cloned()
.collect::<Vec<_>>()
.join("\n")
}
}
pub fn info(&self) -> &VideoInfo {
&self.info
}
pub fn set_cache_capacity(&self, capacity: usize) {
// RGBA preview cap is fixed by memory budget — a user-set 800-frame
// cache at 1080p RGBA would be ~6.6 GB.
if self.format == OutputFormat::Rgba {
return;
}
let clamped = capacity.clamp(MIN_CACHE_CAPACITY, MAX_CACHE_CAPACITY);
let mut st = self.shared.state.lock().unwrap();
st.cache.set_capacity(clamped);
st.lookahead = (clamped as u64) / 2;
}
pub fn decode_at(&mut self, idx: u64) -> Result<Vec<u8>, String> {
match self.decode_at_inner(idx) {
Ok(bytes) => Ok(bytes),
Err(e) => {
// Auto-fallback: any decode failure while hwaccel is on
// retries once with software. Covers both first-frame init
// failures and later seek/codec-edge failures (e.g. a GPU
// path that worked on h264 but chokes on an HEVC segment).
// Once cleared, subsequent failures propagate.
if let Some(prev) = self.hwaccel.take() {
let tail = self.snapshot_stderr_tail();
let notice = format!(
"Hardware decode '{prev}' failed; falling back to software.\nffmpeg said:\n{tail}"
);
eprintln!("[decoder] {notice}");
if let Ok(mut slot) = self.last_notice.lock() {
*slot = Some(notice);
}
self.stop_worker();
self.decode_at_inner(idx)
} else {
Err(e)
}
}
}
}
fn decode_at_inner(&mut self, idx: u64) -> Result<Vec<u8>, String> {
if idx >= self.info.total_frames {
return Err(format!(
"frame {} out of range (total {})",
idx, self.info.total_frames
));
}
// Update reader cursor, check cache, note whether a seek is needed.
let need_seek = {
let mut st = self.shared.state.lock().unwrap();
st.reader_idx = idx;
if let Some(bytes) = st.cache.get(idx) {
self.shared.cond.notify_all();
return Ok(bytes);
}
self.worker.is_none()
|| idx < st.next_idx
|| idx > st.next_idx.saturating_add(SEEK_AHEAD_THRESHOLD)
};
if need_seek {
self.seek_and_start_worker(idx)?;
} else {
self.shared.cond.notify_all();
}
// Wait for the worker to produce idx.
let wait_started = Instant::now();
let mut st = self.shared.state.lock().unwrap();
loop {
if st.cache.contains(idx) {
return st
.cache
.get(idx)
.ok_or_else(|| "frame vanished from cache".to_string());
}
if let Some(e) = &st.error {
return Err(e.clone());
}
if st.eof {
return Err(format!("frame {} not produced (EOF)", idx));
}
if st.shutdown {
return Err("decoder shut down".into());
}
let (next_st, timeout) = self
.shared
.cond
.wait_timeout(st, Duration::from_millis(250))
.unwrap();
st = next_st;
if timeout.timed_out() && wait_started.elapsed() >= DECODE_TIMEOUT {
drop(st);
self.stop_worker();
return Err(format!("timed out decoding frame {idx}"));
}
}
}
fn seek_and_start_worker(&mut self, idx: u64) -> Result<(), String> {
self.stop_worker();
let (child, stdout, stderr) = spawn_ffmpeg(
&self.path,
&self.info,
idx,
self.hwaccel.as_deref(),
self.format,
self.output_dims,
)?;
{
let mut st = self.shared.state.lock().unwrap();
st.stderr_tail.clear();
st.child = Some(child);
st.next_idx = idx;
st.reader_idx = idx;
st.eof = false;
st.error = None;
st.shutdown = false;
}
self.shared.cond.notify_all();
let drain_shared = self.shared.clone();
thread::spawn(move || drain_stderr(stderr, drain_shared));
let format = self.format;
let dims = self.output_dims;
let shared = self.shared.clone();
let handle = thread::spawn(move || {
worker_loop(stdout, shared, format, dims);
});
self.worker = Some(handle);
Ok(())
}
fn stop_worker(&mut self) {
// Take ownership of the child out of shared state and kill it FIRST so
// a pipe read in the worker unblocks with EOF. Then join the worker.
let child_opt = {
let mut st = self.shared.state.lock().unwrap();
st.shutdown = true;
st.child.take()
};
if let Some(mut child) = child_opt {
let _ = child.kill();
let _ = child.wait();
}
self.shared.cond.notify_all();
if let Some(handle) = self.worker.take() {
let _ = handle.join();
}
let mut st = self.shared.state.lock().unwrap();
st.shutdown = false;
st.eof = false;
st.error = None;
}
}
impl Drop for DecoderSession {
fn drop(&mut self) {
self.stop_worker();
}
}
/// Pack a frame for the JS wire: 8-byte little-endian i64 frame index
/// followed by the raw JPEG. The backend is authoritative on which frame
/// actually decoded — JS reads the prefix to bind paints to the real frame
/// number rather than the optimistic request, which matters when codec
/// keyframe-snap or seek-ahead returns something other than the requested
/// index. (Today our ffmpeg path is exact-seek so they always match, but
/// the prefix makes the system robust to that ever changing.)
pub fn pack_frame(idx: u64, jpeg: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(8 + jpeg.len());
out.extend_from_slice(&(idx as i64).to_le_bytes());
out.extend_from_slice(jpeg);
out
}
/// Pick a VAAPI render node. `SAM_TOOL_VAAPI_DEVICE` overrides; otherwise we
/// probe `/dev/dri/renderD128..199` for the first that exists. Multi-GPU
/// systems can end up with the integrated card on D128 and a discrete card
/// on D129+ — the env var is the escape hatch when the auto pick is wrong.
fn vaapi_device() -> String {
if let Ok(p) = std::env::var("SAM_TOOL_VAAPI_DEVICE") {
if !p.is_empty() {
return p;
}
}
for n in 128..200u32 {
let p = format!("/dev/dri/renderD{n}");
if std::path::Path::new(&p).exists() {
return p;
}
}
"/dev/dri/renderD128".to_string()
}
/// Compute the (width, height) ffmpeg will output. Scales down to
/// preview_max_height only when source is taller; preserves aspect via a
/// rounded-to-even width. Returns source dims when no scaling applies.
/// The reader uses these to frame raw video output, so they MUST match
/// what ffmpeg actually emits — we set them explicitly via `-s WxH` in
/// the Rgba path to take the guesswork out.
fn compute_output_dims(info: &VideoInfo, preview_max_height: Option<u32>) -> (u32, u32) {
if let Some(h) = preview_max_height {
if info.height > h && info.height > 0 {
// Maintain aspect; round width to nearest even (ffmpeg requires
// even dims for many filters).
let w = ((info.width as u64 * h as u64) + (info.height as u64 / 2))
/ (info.height as u64);
let w_even = ((w as u32) + 1) & !1;
return (w_even.max(2), h);
}
}
// Even source dims for safety (almost always already even).
let w = (info.width + 1) & !1;
let h = (info.height + 1) & !1;
(w.max(2), h.max(2))
}
fn spawn_ffmpeg(
path: &PathBuf,
info: &VideoInfo,
idx: u64,
hwaccel: Option<&str>,
format: OutputFormat,
out_dims: (u32, u32),
) -> Result<(Child, BufReader<ChildStdout>, ChildStderr), String> {
let path_str = path.to_str().ok_or_else(|| "non-utf8 path".to_string())?;
let t = idx as f64 / info.fps;
let mut cmd = Command::new("ffmpeg");
cmd.args(["-v", "error", "-threads", "0"]);
// -hwaccel must come BEFORE -i. Output goes to CPU memory by default
// (no -hwaccel_output_format), so the downstream mjpeg encoder works
// unchanged. If the GPU init fails, ffmpeg will exit nonzero and the
// worker thread surfaces the error via shared.error.
match hwaccel {
Some("auto") => {
cmd.args(["-hwaccel", "auto"]);
}
Some("vaapi") => {
let dev = vaapi_device();
cmd.args(["-hwaccel", "vaapi", "-vaapi_device", &dev]);
}
Some("cuda") => {
cmd.args(["-hwaccel", "cuda"]);
}
Some("vdpau") => {
cmd.args(["-hwaccel", "vdpau"]);
}
_ => {}
}
// Single fast-seek `-ss` before `-i`. The old `(-ss t-1) -i (-ss 1.0)`
// double-seek pattern decoded ~30 hidden warmup frames per seek; on
// typical content the fast-seek is accurate enough and seeks take
// milliseconds instead of hundreds of ms.
cmd.args(["-ss", &format!("{:.6}", t)]);
cmd.args(["-i", path_str]);
cmd.args(["-an"]);
// Pin output dimensions exactly. For Mjpeg this is cosmetic (matches
// source); for Rgba it MUST match what the reader expects to byte-frame
// raw video output.
let (ow, oh) = out_dims;
if ow != info.width || oh != info.height {
cmd.args(["-s", &format!("{ow}x{oh}")]);
}
match format {
OutputFormat::Mjpeg => {
cmd.args([
"-q:v",
"3",
"-f",
"image2pipe",
"-vcodec",
"mjpeg",
"pipe:1",
]);
}
OutputFormat::Rgba => {
// Raw RGBA, no entropy coding. Worker grabs the bytes verbatim
// into ImageData, eliminating the CPU MJPEG encode (ffmpeg)
// and the CPU JPEG decode (createImageBitmap on WebKit2GTK)
// that bracketed every preview frame in the old pipeline.
cmd.args([
"-pix_fmt",
"rgba",
"-f",
"rawvideo",
"pipe:1",
]);
}
}
// stderr is captured (was /dev/null) into a non-blocking drain thread
// that pushes lines into shared.state.stderr_tail. Bounded buffer, no
// join on EOF — the earlier 150ms wait-for-stderr-tail problem was the
// synchronous WAIT, not the capture itself.
cmd.stdout(Stdio::piped()).stderr(Stdio::piped());
let mut child = cmd.spawn().map_err(|e| format!("ffmpeg spawn: {e}"))?;
let stdout = child
.stdout
.take()
.ok_or_else(|| "ffmpeg stdout unavailable".to_string())?;
let stderr = child
.stderr
.take()
.ok_or_else(|| "ffmpeg stderr unavailable".to_string())?;
Ok((
child,
BufReader::with_capacity(READ_BUF_BYTES, stdout),
stderr,
))
}
/// Read ffmpeg stderr line-by-line into a bounded ring buffer on shared
/// state. Exits when ffmpeg closes the pipe (EOF) or the read errors.
/// Cheap: no blocking on the main thread, used only to surface diagnostic
/// info when something has gone wrong.
fn drain_stderr(stderr: ChildStderr, shared: Arc<Shared>) {
let reader = std::io::BufReader::new(stderr);
for line_res in reader.lines() {
let line = match line_res {
Ok(l) => l,
Err(_) => return,
};
let trimmed = line.trim().to_string();
if trimmed.is_empty() {
continue;
}
let mut st = match shared.state.lock() {
Ok(g) => g,
Err(_) => return,
};
st.stderr_tail.push_back(trimmed);
while st.stderr_tail.len() > STDERR_TAIL_LINES {
st.stderr_tail.pop_front();
}
}
}
fn worker_loop(
mut stdout: BufReader<ChildStdout>,
shared: Arc<Shared>,
format: OutputFormat,
dims: (u32, u32),
) {
// Pre-compute the exact frame size for the Rgba path so we don't pay
// for marker scanning that doesn't exist in raw video.
let raw_frame_bytes = (dims.0 as usize) * (dims.1 as usize) * 4;
loop {
{
let st = shared.state.lock().unwrap();
if st.shutdown {
return;
}
let ahead = st.next_idx.saturating_sub(st.reader_idx);
if ahead >= st.lookahead {
let (st2, _) = shared
.cond
.wait_timeout(st, Duration::from_millis(100))
.unwrap();
if st2.shutdown {
return;
}
continue;
}
}
let read_result = match format {
OutputFormat::Mjpeg => read_jpeg_bytes(&mut stdout),
OutputFormat::Rgba => read_raw_frame(&mut stdout, raw_frame_bytes, dims),
};
match read_result {
Ok(bytes) => {
let mut st = shared.state.lock().unwrap();
if st.shutdown {
return;
}
let idx = st.next_idx;
st.cache.insert(idx, bytes);
st.next_idx = idx + 1;
shared.cond.notify_all();
}
Err(e) => {
let mut st = shared.state.lock().unwrap();
if e.contains("pipe closed") || e.contains("EOF") {
st.eof = true;
} else {
st.error = Some(e);
}
shared.cond.notify_all();
return;
}
}
}
}
/// Read exactly one raw RGBA frame from the pipe and prefix it with
/// `[width:u32 LE, height:u32 LE]`. The JS worker then knows the
/// dimensions without having to track them out-of-band. Returns
/// `Err("pipe closed")` (or "EOF") so the worker loop's EOF detection
/// branch fires cleanly.
fn read_raw_frame(
reader: &mut BufReader<ChildStdout>,
frame_bytes: usize,
dims: (u32, u32),
) -> Result<Vec<u8>, String> {
use std::io::Read;
let mut out = Vec::with_capacity(8 + frame_bytes);
out.extend_from_slice(&dims.0.to_le_bytes());
out.extend_from_slice(&dims.1.to_le_bytes());
out.resize(8 + frame_bytes, 0);
let buf = &mut out[8..];
let mut filled = 0usize;
while filled < frame_bytes {
match reader.read(&mut buf[filled..]) {
Ok(0) => {
return Err(format!(
"ffmpeg pipe closed mid-frame ({filled}/{frame_bytes} bytes)"
));
}
Ok(n) => filled += n,
Err(e) => return Err(format!("pipe read: {e}")),
}
}
Ok(out)
}
/// Read one JPEG (SOI 0xFF 0xD8 .. EOI 0xFF 0xD9) from an image2pipe stream.
///
/// Chunked scan: for each `fill_buf` chunk we vector-scan for the marker
/// byte (`slice::iter().position` compiles to a SIMD memchr in release),
/// then `extend_from_slice` the body wholesale. The previous per-byte
/// `for &b in buf { push(b); ... }` cost ~3 ns/byte and at ~150 KB/frame ×
/// 30 fps it was a measurable chunk of the worker's budget. The bulk-copy
/// path measures around 10-20× faster on warm caches.
fn read_jpeg_bytes(reader: &mut BufReader<ChildStdout>) -> Result<Vec<u8>, String> {
let mut out: Vec<u8> = Vec::with_capacity(64 * 1024);
let mut in_jpeg = false;
// Last byte of the previous chunk; needed because the 2-byte JPEG
// marker (0xFF 0xMM) can straddle a fill_buf boundary.
let mut last: u8 = 0;
loop {
let buf = reader.fill_buf().map_err(|e| format!("pipe read: {e}"))?;
if buf.is_empty() {
return Err("ffmpeg pipe closed before frame completed".into());
}
let n = buf.len();
if !in_jpeg {
match find_marker(buf, last, 0xD8) {
Some(d8_pos) => {
in_jpeg = true;
out.push(0xFF);
out.push(0xD8);
let after = d8_pos + 1;
let rest = &buf[after..];
// Same-chunk EOI is the common case for small frames.
if let Some(d9_pos) = find_marker(rest, 0, 0xD9) {
out.extend_from_slice(&rest[..=d9_pos]);
reader.consume(after + d9_pos + 1);
return Ok(out);
}
out.extend_from_slice(rest);
last = rest.last().copied().unwrap_or(0);
reader.consume(n);
}
None => {
last = buf[n - 1];
reader.consume(n);
}
}
continue;
}
if let Some(d9_pos) = find_marker(buf, last, 0xD9) {
out.extend_from_slice(&buf[..=d9_pos]);
reader.consume(d9_pos + 1);
return Ok(out);
}
out.extend_from_slice(buf);
last = buf[n - 1];
reader.consume(n);
}
}
/// Find the index of `marker` byte preceded by 0xFF within `buf`.
/// `prev` is the last byte of the previous chunk (0 if none) so a marker
/// that straddles a fill_buf boundary is still detected.
fn find_marker(buf: &[u8], prev: u8, marker: u8) -> Option<usize> {
if buf.is_empty() {
return None;
}
if prev == 0xFF && buf[0] == marker {
return Some(0);
}
let mut start = 0usize;
while start + 1 < buf.len() {
// SIMD-friendly scan for 0xFF; LLVM lowers this to a memchr loop.
let off = buf[start..buf.len() - 1]
.iter()
.position(|&b| b == 0xFF)?;
let ff_pos = start + off;
if buf[ff_pos + 1] == marker {
return Some(ff_pos + 1);
}
start = ff_pos + 1;
}
None
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::{BufRead, Cursor, Read, Result as IoResult};
#[test]
fn find_marker_basic() {
// 0xFF 0xD8 at index 2-3
let buf = &[0x00, 0x11, 0xFF, 0xD8, 0x22][..];
assert_eq!(find_marker(buf, 0, 0xD8), Some(3));
}
#[test]
fn find_marker_no_match() {
let buf = &[0xFF, 0xC0, 0xFF, 0xC4, 0x00][..];
assert_eq!(find_marker(buf, 0, 0xD8), None);
}
#[test]
fn find_marker_boundary() {
// prev chunk ended on 0xFF, this chunk starts with 0xD9
let buf = &[0xD9, 0x00, 0x00][..];
assert_eq!(find_marker(buf, 0xFF, 0xD9), Some(0));
}
#[test]
fn find_marker_trailing_ff() {
// 0xFF at the last byte — has no successor in this chunk, so no
// match here; the caller will carry it as `last` to the next chunk.
let buf = &[0x00, 0x11, 0xFF][..];
assert_eq!(find_marker(buf, 0, 0xD9), None);
}
#[test]
fn find_marker_skips_false_positive() {
// 0xFF followed by 0xC0 (not our marker), then real 0xFF 0xD9 later.
let buf = &[0xFF, 0xC0, 0x11, 0xFF, 0xD9, 0x22][..];
assert_eq!(find_marker(buf, 0, 0xD9), Some(4));
}
// ChunkedReader: a BufRead that returns the source one configurable
// chunk at a time, so we can exercise the boundary logic in
// read_jpeg_bytes.
struct ChunkedReader {
data: Vec<u8>,
pos: usize,
chunk: usize,
held: Vec<u8>,
}
impl ChunkedReader {
fn new(data: Vec<u8>, chunk: usize) -> Self {
Self {
data,
pos: 0,
chunk,
held: Vec::new(),
}
}
}
impl Read for ChunkedReader {
fn read(&mut self, _: &mut [u8]) -> IoResult<usize> {
unreachable!("BufRead path only")
}
}
impl BufRead for ChunkedReader {
fn fill_buf(&mut self) -> IoResult<&[u8]> {
if self.held.is_empty() && self.pos < self.data.len() {
let end = (self.pos + self.chunk).min(self.data.len());
self.held = self.data[self.pos..end].to_vec();
self.pos = end;
}
Ok(&self.held)
}
fn consume(&mut self, n: usize) {
self.held.drain(..n);
}
}
// Test-only sibling of read_jpeg_bytes that takes any BufRead. Keeps
// the production signature constrained to BufReader<ChildStdout>.
fn read_jpeg_generic<R: BufRead>(reader: &mut R) -> Result<Vec<u8>, String> {
let mut out: Vec<u8> = Vec::with_capacity(64 * 1024);
let mut in_jpeg = false;
let mut last: u8 = 0;
loop {
let buf = reader.fill_buf().map_err(|e| format!("pipe read: {e}"))?;
if buf.is_empty() {
return Err("ffmpeg pipe closed before frame completed".into());
}
let n = buf.len();
if !in_jpeg {
match find_marker(buf, last, 0xD8) {
Some(d8_pos) => {
in_jpeg = true;
out.push(0xFF);
out.push(0xD8);
let after = d8_pos + 1;
let rest_len = n - after;
// borrow ends after this if-block; copy data out first
let (rest_copy, rest_last) = {
let rest = &buf[after..];
if let Some(d9_pos) = find_marker(rest, 0, 0xD9) {
out.extend_from_slice(&rest[..=d9_pos]);
let total = after + d9_pos + 1;
reader.consume(total);
return Ok(out);
}
(rest.to_vec(), rest.last().copied().unwrap_or(0))
};
out.extend_from_slice(&rest_copy);
last = rest_last;
let _ = rest_len;
reader.consume(n);
}
None => {
last = buf[n - 1];
reader.consume(n);
}
}
continue;
}
if let Some(d9_pos) = find_marker(buf, last, 0xD9) {
let take = d9_pos + 1;
let body: Vec<u8> = buf[..take].to_vec();
out.extend_from_slice(&body);
reader.consume(take);
return Ok(out);
}
let last_byte = buf[n - 1];
let body: Vec<u8> = buf.to_vec();
out.extend_from_slice(&body);
last = last_byte;
reader.consume(n);
}
}
fn make_jpeg(prefix: &[u8], body: &[u8], suffix: &[u8]) -> Vec<u8> {
let mut v = Vec::new();
v.extend_from_slice(prefix);
v.extend_from_slice(&[0xFF, 0xD8]);
v.extend_from_slice(body);
v.extend_from_slice(&[0xFF, 0xD9]);
v.extend_from_slice(suffix);
v
}
fn expected_jpeg(body: &[u8]) -> Vec<u8> {
let mut v = vec![0xFF, 0xD8];
v.extend_from_slice(body);
v.extend_from_slice(&[0xFF, 0xD9]);
v
}
#[test]
fn read_jpeg_single_chunk() {
let body = vec![0x11, 0x22, 0x33, 0x44];
let stream = make_jpeg(&[0xAA], &body, &[0xBB]);
let mut r = ChunkedReader::new(stream, 4096);
let got = read_jpeg_generic(&mut r).unwrap();
assert_eq!(got, expected_jpeg(&body));
}
#[test]
fn read_jpeg_byte_at_a_time() {
let body = vec![0x10, 0x20, 0xFF, 0x00, 0x30, 0x40];
let stream = make_jpeg(&[0x99, 0xFF, 0xC0], &body, &[]);
let mut r = ChunkedReader::new(stream, 1);
let got = read_jpeg_generic(&mut r).unwrap();
assert_eq!(got, expected_jpeg(&body));
}
#[test]
fn read_jpeg_soi_straddles_chunks() {
// Force a chunk boundary that ends with 0xFF and the next starts
// with 0xD8 (the SOI's second byte).
let body = vec![0xAA, 0xBB];
let mut stream = vec![0xFF]; // chunk1 = [0xFF]
stream.extend_from_slice(&[0xD8]); // chunk2 begins with 0xD8
stream.extend_from_slice(&body);
stream.extend_from_slice(&[0xFF, 0xD9]);
let mut r = ChunkedReader::new(stream, 1);
let got = read_jpeg_generic(&mut r).unwrap();
assert_eq!(got, expected_jpeg(&body));
}
#[test]
fn read_jpeg_eoi_straddles_chunks() {
// chunk1 ends on the 0xFF of EOI, chunk2 starts with 0xD9.
let body = vec![0x10, 0x20, 0x30];
let mut stream = vec![0xFF, 0xD8];
stream.extend_from_slice(&body);
stream.push(0xFF); // chunk boundary right after this
stream.push(0xD9);
let mut r = ChunkedReader::new(stream, 6); // body+SOI = 5, then 0xFF lands at idx 5 → fits in first chunk; need a chunk size that splits between 0xFF and 0xD9. SOI(2)+body(3)+0xFF = 6 bytes; chunk size 6 → first chunk includes through 0xFF. Then chunk2 = [0xD9].
let got = read_jpeg_generic(&mut r).unwrap();
assert_eq!(got, expected_jpeg(&body));
}
#[test]
fn read_jpeg_false_eoi_in_body() {
// 0xFF 0xC0 (a JFIF marker) inside body should not be mistaken for EOI.
let body = vec![0x11, 0xFF, 0xC0, 0x22, 0xFF, 0xDB, 0x33];
let stream = make_jpeg(&[], &body, &[]);
let mut r = ChunkedReader::new(stream, 3);
let got = read_jpeg_generic(&mut r).unwrap();
assert_eq!(got, expected_jpeg(&body));
}
#[test]
fn fifo_cache_evicts_oldest_in_order() {
let mut c = FrameCache::new(3);
c.insert(1, vec![1]);
c.insert(2, vec![2]);
c.insert(3, vec![3]);
c.insert(4, vec![4]); // evicts 1
assert!(!c.contains(1));
assert_eq!(c.get(2), Some(vec![2]));
assert_eq!(c.get(3), Some(vec![3]));
assert_eq!(c.get(4), Some(vec![4]));
}
#[test]
fn fifo_cache_overwrite_keeps_order() {
let mut c = FrameCache::new(2);
c.insert(1, vec![1]);
c.insert(2, vec![2]);
c.insert(1, vec![10]); // overwrite, do NOT reorder
c.insert(3, vec![3]); // evicts 1 (still the oldest in queue)
assert!(!c.contains(1));
assert_eq!(c.get(2), Some(vec![2]));
assert_eq!(c.get(3), Some(vec![3]));
}
#[test]
fn cache_used() {
// Suppress unused-import warning under cfg(test).
let _ = Cursor::new(vec![1u8]);
}
}