528 lines
18 KiB
Rust
528 lines
18 KiB
Rust
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};
|
||
|
||
pub const DEFAULT_CACHE_CAPACITY: usize = 160;
|
||
pub const MIN_CACHE_CAPACITY: usize = 50;
|
||
pub const MAX_CACHE_CAPACITY: usize = 800;
|
||
const READ_BUF_BYTES: usize = 256 * 1024;
|
||
/// If the requested frame is this far ahead of the worker, kill + reseek.
|
||
const SEEK_AHEAD_THRESHOLD: u64 = 60;
|
||
const DECODE_TIMEOUT: Duration = Duration::from_secs(15);
|
||
/// Keep the last N stderr lines from ffmpeg so error messages can include
|
||
/// the real reason (e.g. "vaapi: failed to load driver") instead of just
|
||
/// "pipe closed before frame completed".
|
||
const STDERR_TAIL_LINES: usize = 32;
|
||
|
||
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)
|
||
}
|
||
|
||
fn get(&mut self, idx: u64) -> Option<Vec<u8>> {
|
||
let bytes = self.map.get(&idx)?.clone();
|
||
if let Some(pos) = self.order.iter().position(|&i| i == idx) {
|
||
self.order.remove(pos);
|
||
}
|
||
self.order.push_back(idx);
|
||
Some(bytes)
|
||
}
|
||
|
||
fn insert(&mut self, idx: u64, bytes: Vec<u8>) {
|
||
if self.map.contains_key(&idx) {
|
||
if let Some(pos) = self.order.iter().position(|&i| i == idx) {
|
||
self.order.remove(pos);
|
||
}
|
||
} else 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,
|
||
}
|
||
|
||
struct Shared {
|
||
state: Mutex<SharedState>,
|
||
cond: Condvar,
|
||
/// Ring buffer of recent ffmpeg stderr lines, filled by a sibling drain
|
||
/// thread. Cleared at each new seek.
|
||
stderr_tail: Mutex<VecDeque<String>>,
|
||
}
|
||
|
||
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>,
|
||
shared: Arc<Shared>,
|
||
worker: Option<JoinHandle<()>>,
|
||
/// Sibling thread draining ffmpeg's stderr into `shared.stderr_tail`.
|
||
/// Joined in `stop_worker` so the tail doesn't get clobbered by a
|
||
/// straggling drain across seeks.
|
||
stderr_drain: Option<JoinHandle<()>>,
|
||
}
|
||
|
||
impl DecoderSession {
|
||
pub fn open(path: PathBuf, info: VideoInfo, hwaccel: Option<String>) -> Self {
|
||
Self {
|
||
path,
|
||
info,
|
||
hwaccel: hwaccel.filter(|s| !s.is_empty()),
|
||
shared: Arc::new(Shared {
|
||
state: Mutex::new(SharedState {
|
||
child: None,
|
||
cache: FrameCache::new(DEFAULT_CACHE_CAPACITY),
|
||
next_idx: 0,
|
||
reader_idx: 0,
|
||
lookahead: (DEFAULT_CACHE_CAPACITY as u64) / 2,
|
||
eof: false,
|
||
error: None,
|
||
shutdown: false,
|
||
}),
|
||
cond: Condvar::new(),
|
||
stderr_tail: Mutex::new(VecDeque::new()),
|
||
}),
|
||
worker: None,
|
||
stderr_drain: None,
|
||
}
|
||
}
|
||
|
||
pub fn info(&self) -> &VideoInfo {
|
||
&self.info
|
||
}
|
||
|
||
pub fn set_cache_capacity(&self, capacity: usize) {
|
||
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() {
|
||
eprintln!(
|
||
"[decoder] hwaccel '{prev}' failed ({e}); falling back to software"
|
||
);
|
||
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())?;
|
||
|
||
{
|
||
let mut st = self.shared.state.lock().unwrap();
|
||
st.child = Some(child);
|
||
st.next_idx = idx;
|
||
st.reader_idx = idx;
|
||
st.eof = false;
|
||
st.error = None;
|
||
st.shutdown = false;
|
||
}
|
||
self.shared.stderr_tail.lock().unwrap().clear();
|
||
self.shared.cond.notify_all();
|
||
|
||
// Drain ffmpeg's stderr into a ring buffer so error messages can
|
||
// include the real reason. We must drain even on success — leaving
|
||
// stderr unread would let the pipe buffer fill and stall ffmpeg.
|
||
let drain_shared = self.shared.clone();
|
||
self.stderr_drain = Some(thread::spawn(move || {
|
||
let reader = BufReader::new(stderr);
|
||
for line in reader.lines().map_while(Result::ok) {
|
||
let mut tail = drain_shared.stderr_tail.lock().unwrap();
|
||
if tail.len() >= STDERR_TAIL_LINES {
|
||
tail.pop_front();
|
||
}
|
||
tail.push_back(line);
|
||
}
|
||
}));
|
||
|
||
let shared = self.shared.clone();
|
||
let handle = thread::spawn(move || {
|
||
worker_loop(stdout, shared);
|
||
});
|
||
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();
|
||
}
|
||
// Killing the child closed stderr, so the drain thread sees EOF and
|
||
// exits — joining it here prevents a straggling write into the new
|
||
// session's stderr_tail after the next seek clears it.
|
||
if let Some(handle) = self.stderr_drain.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();
|
||
}
|
||
}
|
||
|
||
/// 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()
|
||
}
|
||
|
||
fn spawn_ffmpeg(
|
||
path: &PathBuf,
|
||
info: &VideoInfo,
|
||
idx: u64,
|
||
hwaccel: Option<&str>,
|
||
) -> 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"]);
|
||
}
|
||
_ => {}
|
||
}
|
||
if t > 1.0 {
|
||
cmd.args(["-ss", &format!("{:.3}", t - 1.0)]);
|
||
cmd.args(["-i", path_str]);
|
||
cmd.args(["-ss", "1.0"]);
|
||
} else {
|
||
cmd.args(["-i", path_str]);
|
||
cmd.args(["-ss", &format!("{:.3}", t)]);
|
||
}
|
||
cmd.args([
|
||
"-an",
|
||
"-q:v",
|
||
"3",
|
||
"-f",
|
||
"image2pipe",
|
||
"-vcodec",
|
||
"mjpeg",
|
||
"pipe:1",
|
||
]);
|
||
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,
|
||
))
|
||
}
|
||
|
||
fn worker_loop(mut stdout: BufReader<ChildStdout>, shared: Arc<Shared>) {
|
||
loop {
|
||
// Backpressure: wait if we're already far enough ahead.
|
||
{
|
||
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;
|
||
}
|
||
}
|
||
|
||
match read_jpeg_bytes(&mut stdout) {
|
||
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) => {
|
||
// ffmpeg closes stderr right alongside stdout; on real
|
||
// failures the drain thread may not have flushed yet, so
|
||
// we wait briefly. On a clean exit there's nothing to
|
||
// flush — collect_stderr_tail short-circuits via the
|
||
// child's exit code so end-of-stream EOF isn't penalised.
|
||
let tail = collect_stderr_tail(&shared, Duration::from_millis(150));
|
||
let mut st = shared.state.lock().unwrap();
|
||
if !tail.is_empty() {
|
||
// With `-v error`, anything on stderr means ffmpeg
|
||
// actually failed (vs a clean exit) — surface it as a
|
||
// real error so decode_at can fall back / report.
|
||
st.error = Some(format!("{e}: ffmpeg: {tail}"));
|
||
} else if e.contains("pipe closed") || e.contains("EOF") {
|
||
st.eof = true;
|
||
} else {
|
||
st.error = Some(e);
|
||
}
|
||
shared.cond.notify_all();
|
||
return;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
fn read_tail_now(shared: &Arc<Shared>) -> Option<String> {
|
||
let tail = shared.stderr_tail.lock().unwrap();
|
||
if tail.is_empty() {
|
||
None
|
||
} else {
|
||
Some(tail.iter().cloned().collect::<Vec<_>>().join(" | "))
|
||
}
|
||
}
|
||
|
||
fn collect_stderr_tail(shared: &Arc<Shared>, max_wait: Duration) -> String {
|
||
// Fast path: drain already flushed — common when ffmpeg writes its
|
||
// error before exiting and stdout EOF arrives after the drain.
|
||
if let Some(tail) = read_tail_now(shared) {
|
||
return tail;
|
||
}
|
||
// Skip the poll on clean ffmpeg exit (exit code 0): natural end-of-
|
||
// stream has nothing on stderr (we run with `-v error`), and 150ms ×
|
||
// every EOF seek near the final frames adds up.
|
||
let exited_cleanly = {
|
||
let mut st = shared.state.lock().unwrap();
|
||
st.child
|
||
.as_mut()
|
||
.and_then(|c| c.try_wait().ok().flatten())
|
||
.map(|status| status.success())
|
||
.unwrap_or(false)
|
||
};
|
||
if exited_cleanly {
|
||
return String::new();
|
||
}
|
||
// Slow path: ffmpeg exited non-zero (or hasn't exited yet) and the
|
||
// drain thread hasn't flushed. Poll briefly so the error message
|
||
// carries the real reason.
|
||
let deadline = Instant::now() + max_wait;
|
||
loop {
|
||
if let Some(tail) = read_tail_now(shared) {
|
||
return tail;
|
||
}
|
||
if Instant::now() >= deadline {
|
||
return String::new();
|
||
}
|
||
thread::sleep(Duration::from_millis(10));
|
||
}
|
||
}
|
||
|
||
fn read_jpeg_bytes(reader: &mut BufReader<ChildStdout>) -> Result<Vec<u8>, String> {
|
||
let mut out = Vec::with_capacity(64 * 1024);
|
||
let mut last = 0u8;
|
||
let mut in_jpeg = false;
|
||
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 mut consumed = 0usize;
|
||
for &b in buf {
|
||
consumed += 1;
|
||
if !in_jpeg {
|
||
if last == 0xFF && b == 0xD8 {
|
||
in_jpeg = true;
|
||
out.push(0xFF);
|
||
out.push(0xD8);
|
||
}
|
||
last = b;
|
||
} else {
|
||
out.push(b);
|
||
if last == 0xFF && b == 0xD9 {
|
||
reader.consume(consumed);
|
||
return Ok(out);
|
||
}
|
||
last = b;
|
||
}
|
||
}
|
||
reader.consume(consumed);
|
||
}
|
||
}
|