Skip to main content

pito_codecs/media/front/
ffmpeg.rs

1use std::collections::{BTreeMap, HashMap};
2use std::io::Read;
3use std::path::{Path, PathBuf};
4use std::process::{ChildStdin, ChildStdout, Command, Stdio};
5use std::sync::{Mutex, OnceLock, PoisonError};
6use std::thread::JoinHandle;
7
8use crate::media::{Codec, Direction, Error, PixelFormat, Result};
9
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub struct Listed {
12    pub name: String,
13    pub codec: String,
14    pub video: bool,
15    pub hardware: bool,
16    pub description: String,
17}
18
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct Ffmpeg {
21    pub path: PathBuf,
22    pub ffprobe: Option<PathBuf>,
23    pub version: String,
24    pub encoders: Vec<Listed>,
25    pub decoders: Vec<Listed>,
26    pub hwaccels: Vec<String>,
27}
28
29const HARDWARE: [&str; 11] = [
30    "_vaapi",
31    "_nvenc",
32    "_cuvid",
33    "_qsv",
34    "_videotoolbox",
35    "_mf",
36    "_amf",
37    "_v4l2m2m",
38    "_vulkan",
39    "_d3d12va",
40    "_mediacodec",
41];
42
43fn cache() -> &'static Mutex<HashMap<PathBuf, Ffmpeg>> {
44    static CACHE: OnceLock<Mutex<HashMap<PathBuf, Ffmpeg>>> = OnceLock::new();
45    CACHE.get_or_init(|| Mutex::new(HashMap::new()))
46}
47
48impl Ffmpeg {
49    pub fn find(user: Option<&Path>) -> Result<Ffmpeg> {
50        let path = match user {
51            Some(p) if p.is_file() => p.to_path_buf(),
52            Some(p) => {
53                return Err(Error::Unsupported(format!("no ffmpeg at {}", p.display())));
54            }
55            None => on_path(exe("ffmpeg"))
56                .ok_or_else(|| Error::Unsupported("no ffmpeg on PATH".to_string()))?,
57        };
58        if let Some(f) = cache()
59            .lock()
60            .unwrap_or_else(PoisonError::into_inner)
61            .get(&path)
62        {
63            return Ok(f.clone());
64        }
65        let f = Ffmpeg::load(path.clone())?;
66        cache()
67            .lock()
68            .unwrap_or_else(PoisonError::into_inner)
69            .insert(path, f.clone());
70        Ok(f)
71    }
72
73    fn load(path: PathBuf) -> Result<Ffmpeg> {
74        let ffprobe = path
75            .parent()
76            .map(|d| d.join(exe("ffprobe")))
77            .filter(|p| p.is_file())
78            .or_else(|| on_path(exe("ffprobe")));
79        let mut f = Ffmpeg {
80            path,
81            ffprobe,
82            version: String::new(),
83            encoders: Vec::new(),
84            decoders: Vec::new(),
85            hwaccels: Vec::new(),
86        };
87        let (version, encoders, decoders, hwaccels) = std::thread::scope(|s| {
88            let fr = &f;
89            let e = s.spawn(move || fr.run(&["-encoders"]));
90            let d = s.spawn(move || fr.run(&["-decoders"]));
91            let h = s.spawn(move || fr.run(&["-hwaccels"]));
92            let v = fr.run(&["-version"]);
93            let join = |r: std::thread::Result<Result<Vec<u8>>>| {
94                r.unwrap_or_else(|_| Err(Error::Backend("ffmpeg listing failed".to_string())))
95            };
96            (v, join(e.join()), join(d.join()), join(h.join()))
97        });
98        f.version = String::from_utf8_lossy(&version?)
99            .lines()
100            .next()
101            .unwrap_or("")
102            .to_string();
103        f.encoders = listing(&String::from_utf8_lossy(&encoders?));
104        f.decoders = listing(&String::from_utf8_lossy(&decoders?));
105        f.hwaccels = String::from_utf8_lossy(&hwaccels?)
106            .lines()
107            .skip_while(|l| !l.starts_with("Hardware acceleration methods"))
108            .skip(1)
109            .map(str::trim)
110            .filter(|l| !l.is_empty())
111            .map(str::to_string)
112            .collect();
113        Ok(f)
114    }
115
116    pub fn can(&self, codec: Codec, dir: Direction) -> bool {
117        let list = match dir {
118            Direction::Decode => &self.decoders,
119            Direction::Encode => &self.encoders,
120        };
121        list.iter().any(|l| l.codec == codec_name(codec))
122    }
123
124    pub fn encoder_for(&self, codec: Codec, picked: Option<&str>) -> Option<String> {
125        let wanted = codec_name(codec);
126        let ours = |n: &str| {
127            self.encoders
128                .iter()
129                .any(|l| l.name == n && l.codec == wanted)
130        };
131        if let Some(p) = picked.filter(|p| ours(p)) {
132            return Some(p.to_string());
133        }
134        let order: &[&str] = match codec {
135            Codec::H264 => &["libx264", "h264_vaapi", "h264_nvenc", "h264_videotoolbox"],
136            Codec::Hevc => &["libx265", "hevc_vaapi", "hevc_nvenc", "hevc_videotoolbox"],
137            Codec::Aac => &["aac", "libfdk_aac"],
138            Codec::Vp9 => &["libvpx-vp9", "vp9_vaapi"],
139            Codec::Opus => &["libopus", "opus"],
140            Codec::Ac3 => &["ac3"],
141            Codec::Eac3 => &["eac3"],
142            Codec::Mp3 => &["libmp3lame"],
143        };
144        order
145            .iter()
146            .find(|n| ours(n))
147            .map(|n| n.to_string())
148            .or_else(|| {
149                self.encoders
150                    .iter()
151                    .find(|l| l.codec == wanted)
152                    .map(|l| l.name.clone())
153            })
154    }
155
156    pub fn decoder_args(&self, picked: Option<&str>) -> Vec<String> {
157        match picked {
158            Some(p) if self.hwaccels.iter().any(|h| h == p) => {
159                vec!["-hwaccel".into(), p.into()]
160            }
161            Some(p) if self.decoders.iter().any(|l| l.name == p) => {
162                vec!["-c:v".into(), p.into()]
163            }
164            _ => Vec::new(),
165        }
166    }
167
168    pub fn command(&self) -> Command {
169        let mut c = Command::new(&self.path);
170        c.args(["-nostdin", "-hide_banner", "-v", "error"]);
171        background(&mut c);
172        c
173    }
174
175    pub fn run(&self, args: &[&str]) -> Result<Vec<u8>> {
176        let mut c = Command::new(&self.path);
177        c.args(["-hide_banner"]).args(args);
178        background(&mut c);
179        output(c, "ffmpeg")
180    }
181
182    pub fn run_args(&self, args: &[String]) -> Result<Vec<u8>> {
183        let mut c = self.command();
184        c.arg("-y").args(args);
185        output(c, "ffmpeg")
186    }
187
188    fn ffprobe(&self, of: &str, args: &[&str], path: &Path) -> Result<Vec<u8>> {
189        let probe = self.ffprobe.as_ref().ok_or_else(|| {
190            Error::Unsupported(format!("no ffprobe next to {}", self.path.display()))
191        })?;
192        let mut c = Command::new(probe);
193        c.args(["-v", "error", "-of", of]).args(args).arg(path);
194        background(&mut c);
195        output(c, "ffprobe")
196    }
197
198    pub fn probe_flat(&self, args: &[&str], path: &Path) -> Result<BTreeMap<String, String>> {
199        Ok(flat(&String::from_utf8_lossy(
200            &self.ffprobe("flat", args, path)?,
201        )))
202    }
203
204    pub fn probe_csv(&self, args: &[&str], path: &Path) -> Result<String> {
205        Ok(String::from_utf8_lossy(&self.ffprobe("csv=p=0", args, path)?).into_owned())
206    }
207
208    pub fn spawn(&self, args: &[String], stdin: bool, stdout: bool) -> Result<Child> {
209        let mut c = self.command();
210        c.arg("-y")
211            .args(args)
212            .stdin(if stdin { Stdio::piped() } else { Stdio::null() })
213            .stdout(if stdout {
214                Stdio::piped()
215            } else {
216                Stdio::null()
217            })
218            .stderr(Stdio::piped());
219        let mut child = c
220            .spawn()
221            .map_err(|e| Error::Backend(format!("ffmpeg: {e}")))?;
222        let err = child.stderr.take().map(|mut s| {
223            std::thread::spawn(move || {
224                let mut buf = Vec::new();
225                let _ = s.read_to_end(&mut buf);
226                buf
227            })
228        });
229        Ok(Child {
230            stdin: child.stdin.take(),
231            stdout: child.stdout.take(),
232            child,
233            err,
234        })
235    }
236}
237
238pub struct Child {
239    child: std::process::Child,
240    pub stdin: Option<ChildStdin>,
241    pub stdout: Option<ChildStdout>,
242    err: Option<JoinHandle<Vec<u8>>>,
243}
244
245impl Child {
246    pub fn finish(mut self) -> Result<()> {
247        self.stdin.take();
248        self.stdout.take();
249        let status = self.child.wait()?;
250        let err = self
251            .err
252            .take()
253            .and_then(|h| h.join().ok())
254            .unwrap_or_default();
255        if status.success() {
256            Ok(())
257        } else {
258            Err(Error::Backend(format!("ffmpeg: {}", last_line(&err))))
259        }
260    }
261
262    pub fn stderr_now(&mut self) -> String {
263        let _ = self.child.kill();
264        let _ = self.child.wait();
265        self.err
266            .take()
267            .and_then(|h| h.join().ok())
268            .and_then(|e| final_line(&e))
269            .unwrap_or_default()
270    }
271}
272
273impl Drop for Child {
274    fn drop(&mut self) {
275        if let Ok(None) = self.child.try_wait() {
276            let _ = self.child.kill();
277            let _ = self.child.wait();
278        }
279    }
280}
281
282pub fn codec_name(codec: Codec) -> &'static str {
283    match codec {
284        Codec::H264 => "h264",
285        Codec::Hevc => "hevc",
286        Codec::Aac => "aac",
287        Codec::Vp9 => "vp9",
288        Codec::Opus => "opus",
289        Codec::Ac3 => "ac3",
290        Codec::Eac3 => "eac3",
291        Codec::Mp3 => "mp3",
292    }
293}
294
295pub fn codec_of(name: &str) -> Option<Codec> {
296    [
297        Codec::H264,
298        Codec::Hevc,
299        Codec::Aac,
300        Codec::Vp9,
301        Codec::Opus,
302        Codec::Ac3,
303        Codec::Eac3,
304        Codec::Mp3,
305    ]
306    .into_iter()
307    .find(|&c| codec_name(c) == name)
308}
309
310pub fn pix_fmt(format: PixelFormat) -> &'static str {
311    match format {
312        PixelFormat::Yuv420p8 => "yuv420p",
313        PixelFormat::Yuv420p10 => "yuv420p10le",
314    }
315}
316
317fn output(mut c: Command, what: &str) -> Result<Vec<u8>> {
318    let out = c
319        .stdin(Stdio::null())
320        .output()
321        .map_err(|e| Error::Backend(format!("{what}: {e}")))?;
322    if !out.status.success() {
323        return Err(Error::Backend(format!(
324            "{what}: {}",
325            last_line(&out.stderr)
326        )));
327    }
328    Ok(out.stdout)
329}
330
331fn final_line(err: &[u8]) -> Option<String> {
332    String::from_utf8_lossy(err)
333        .lines()
334        .map(str::trim)
335        .rfind(|l| !l.is_empty())
336        .map(str::to_string)
337}
338
339fn last_line(err: &[u8]) -> String {
340    final_line(err).unwrap_or_else(|| "failed with no message".into())
341}
342
343fn exe(name: &str) -> String {
344    if cfg!(windows) {
345        format!("{name}.exe")
346    } else {
347        name.to_string()
348    }
349}
350
351fn on_path(name: String) -> Option<PathBuf> {
352    std::env::split_paths(&std::env::var_os("PATH")?)
353        .map(|d| d.join(&name))
354        .find(|p| p.is_file())
355}
356
357#[cfg(windows)]
358fn background(c: &mut Command) {
359    use std::os::windows::process::CommandExt;
360    c.creation_flags(0x0800_0000);
361}
362
363#[cfg(not(windows))]
364fn background(_c: &mut Command) {}
365
366pub fn listing(text: &str) -> Vec<Listed> {
367    text.lines()
368        .skip_while(|l| !l.trim_start().starts_with("------"))
369        .skip(1)
370        .filter_map(|l| {
371            let mut it = l.split_whitespace();
372            let flags = it.next()?;
373            let name = it.next()?.to_string();
374            let description = it.collect::<Vec<_>>().join(" ");
375            let codec = description
376                .rsplit_once("(codec ")
377                .and_then(|(_, c)| c.strip_suffix(')'))
378                .unwrap_or(&name)
379                .to_string();
380            Some(Listed {
381                hardware: HARDWARE.iter().any(|s| name.ends_with(s)),
382                video: flags.starts_with('V'),
383                name,
384                codec,
385                description,
386            })
387        })
388        .collect()
389}
390
391pub fn flat(text: &str) -> BTreeMap<String, String> {
392    text.lines()
393        .filter_map(|l| {
394            let (k, v) = l.split_once('=')?;
395            let v = v.trim();
396            let v = match v.strip_prefix('"').and_then(|v| v.strip_suffix('"')) {
397                Some(q) => unescape(q),
398                None => v.to_string(),
399            };
400            Some((k.trim().to_string(), v))
401        })
402        .collect()
403}
404
405fn unescape(s: &str) -> String {
406    let mut out = String::with_capacity(s.len());
407    let mut chars = s.chars();
408    while let Some(c) = chars.next() {
409        if c != '\\' {
410            out.push(c);
411            continue;
412        }
413        match chars.next() {
414            Some('n') => out.push('\n'),
415            Some('t') => out.push('\t'),
416            Some('r') => out.push('\r'),
417            Some(o) => out.push(o),
418            None => {}
419        }
420    }
421    out
422}
423
424#[cfg(test)]
425mod tests {
426    use super::*;
427
428    #[test]
429    fn listings_name_the_codec_and_spot_hardware() {
430        let text = " V....D = Video\n ------\n V....D libx264              libx264 H.264 / AVC / MPEG-4 AVC / MPEG-4 part 10 (codec h264)\n V....D h264_vaapi           H.264/AVC (VAAPI) (codec h264)\n A....D aac                  AAC (Advanced Audio Coding)\n";
431        let l = listing(text);
432        assert_eq!(l.len(), 3);
433        assert_eq!((l[0].codec.as_str(), l[0].hardware), ("h264", false));
434        assert_eq!((l[1].codec.as_str(), l[1].hardware), ("h264", true));
435        assert_eq!((l[2].codec.as_str(), l[2].video), ("aac", false));
436    }
437}