first version
This commit is contained in:
351
sam-tool-tauri/src-tauri/src/decoder.rs
Normal file
351
sam-tool-tauri/src-tauri/src/decoder.rs
Normal file
@@ -0,0 +1,351 @@
|
||||
use crate::video::VideoInfo;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::path::PathBuf;
|
||||
use std::process::{Child, ChildStdout, Command, Stdio};
|
||||
use std::sync::{Arc, Condvar, Mutex};
|
||||
use std::thread::{self, JoinHandle};
|
||||
use std::time::Duration;
|
||||
|
||||
pub const DEFAULT_CACHE_CAPACITY: usize = 500;
|
||||
pub const MIN_CACHE_CAPACITY: usize = 200;
|
||||
pub const MAX_CACHE_CAPACITY: usize = 2000;
|
||||
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;
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
pub struct DecoderSession {
|
||||
path: PathBuf,
|
||||
info: VideoInfo,
|
||||
shared: Arc<Shared>,
|
||||
worker: Option<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl DecoderSession {
|
||||
pub fn open(path: PathBuf, info: VideoInfo) -> Self {
|
||||
Self {
|
||||
path,
|
||||
info,
|
||||
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(),
|
||||
}),
|
||||
worker: 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> {
|
||||
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 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());
|
||||
}
|
||||
st = self.shared.cond.wait(st).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
fn seek_and_start_worker(&mut self, idx: u64) -> Result<(), String> {
|
||||
self.stop_worker();
|
||||
|
||||
let (child, stdout) = spawn_ffmpeg(&self.path, &self.info, idx)?;
|
||||
|
||||
{
|
||||
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.cond.notify_all();
|
||||
|
||||
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();
|
||||
}
|
||||
|
||||
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();
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_ffmpeg(
|
||||
path: &PathBuf,
|
||||
info: &VideoInfo,
|
||||
idx: u64,
|
||||
) -> Result<(Child, BufReader<ChildStdout>), 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"]);
|
||||
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())?;
|
||||
Ok((child, BufReader::with_capacity(READ_BUF_BYTES, stdout)))
|
||||
}
|
||||
|
||||
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) => {
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user