Skip to main content

viser_ffmpeg/
encode.rs

1use std::path::{Path, PathBuf};
2use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
3use tokio::io::AsyncBufReadExt;
4use tokio::process::Command;
5
6use crate::{Codec, RateControlMode, Resolution, ffmpeg_path, probe};
7
8/// Parameters for a single encode.
9#[derive(Debug, Clone)]
10pub struct EncodeJob {
11    pub input: String,
12    pub output: String,
13    pub resolution: Option<Resolution>,
14    pub codec: Codec,
15    pub crf: i32,
16    pub rate_control: RateControlMode,
17    pub target_bitrate: f64, // kbps, used for VBR mode
18    pub max_bitrate: f64,    // kbps, used for capped CRF mode
19    pub bufsize: f64,        // kbps, used for capped CRF mode
20    pub preset: String,
21    pub extra_args: Vec<String>,
22}
23
24/// Output of a completed encode.
25#[derive(Debug, Clone)]
26pub struct EncodeResult {
27    pub job: EncodeJob,
28    pub bitrate: f64,       // kbps (average)
29    pub file_size: u64,     // bytes
30    pub duration: Duration, // wall-clock encode time
31}
32
33/// Real-time encoding progress info parsed from FFmpeg.
34#[derive(Debug, Clone, Default)]
35pub struct Progress {
36    pub frame: i64,
37    pub fps: f64,
38    pub bitrate: f64, // kbps
39    pub speed: f64,   // e.g. 2.5x
40    pub time: Duration,
41}
42
43/// Runs an FFmpeg encode job. Progress updates are sent on the channel if provided.
44pub async fn encode(
45    job: EncodeJob,
46    progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
47) -> anyhow::Result<EncodeResult> {
48    match job.rate_control {
49        RateControlMode::Vbr => encode_two_pass(job, progress_tx).await,
50        _ => encode_single_pass(job, progress_tx).await,
51    }
52}
53
54async fn encode_single_pass(
55    job: EncodeJob,
56    progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
57) -> anyhow::Result<EncodeResult> {
58    let args = build_encode_args(&job, EncodePass::Single)?;
59    run_encode(job, args, progress_tx).await
60}
61
62async fn encode_two_pass(
63    job: EncodeJob,
64    progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
65) -> anyhow::Result<EncodeResult> {
66    if job.target_bitrate <= 0.0 {
67        anyhow::bail!("target bitrate must be greater than zero for VBR mode");
68    }
69
70    let passlog_prefix = make_passlog_prefix(&job.output);
71    let cleanup = PasslogCleanup::new(passlog_prefix.clone());
72
73    let first_pass_args = build_encode_args(&job, EncodePass::First(&passlog_prefix))?;
74    run_ffmpeg(first_pass_args, None).await?;
75
76    let second_pass_args = build_encode_args(&job, EncodePass::Second(&passlog_prefix))?;
77    let result = run_encode(job, second_pass_args, progress_tx).await;
78
79    cleanup.run();
80    result
81}
82
83async fn run_encode(
84    job: EncodeJob,
85    args: Vec<String>,
86    progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
87) -> anyhow::Result<EncodeResult> {
88    let start = Instant::now();
89    run_ffmpeg(args, progress_tx).await?;
90
91    let elapsed = start.elapsed();
92
93    // Probe the output to get actual bitrate and file size
94    let meta = std::fs::metadata(&job.output)
95        .map_err(|e| anyhow::anyhow!("failed to stat output: {e}"))?;
96
97    let probe_result = probe(&job.output).await?;
98    let bitrate = probe_result.format.bit_rate as f64 / 1000.0;
99
100    Ok(EncodeResult { job, bitrate, file_size: meta.len(), duration: elapsed })
101}
102
103async fn run_ffmpeg(
104    args: Vec<String>,
105    progress_tx: Option<tokio::sync::mpsc::Sender<Progress>>,
106) -> anyhow::Result<()> {
107    let mut cmd = Command::new(ffmpeg_path());
108    cmd.args(&args).stdout(std::process::Stdio::piped()).stderr(std::process::Stdio::piped());
109
110    let mut child = cmd.spawn().map_err(|e| anyhow::anyhow!("failed to start ffmpeg: {e}"))?;
111
112    // Parse progress from stdout
113    if let Some(stdout) = child.stdout.take() {
114        let tx = progress_tx.clone();
115        tokio::spawn(async move {
116            let reader = tokio::io::BufReader::new(stdout);
117            let mut lines = reader.lines();
118            let mut p = Progress::default();
119            while let Ok(Some(line)) = lines.next_line().await {
120                if parse_progress_line(&line, &mut p) {
121                    if let Some(ref tx) = tx {
122                        let _ = tx.try_send(p.clone());
123                    }
124                }
125            }
126        });
127    }
128
129    let output = child.wait_with_output().await?;
130    if !output.status.success() {
131        let stderr = String::from_utf8_lossy(&output.stderr);
132        anyhow::bail!("ffmpeg encode failed: {stderr}");
133    }
134
135    Ok(())
136}
137
138/// Copies a segment of a video file without re-encoding.
139pub async fn extract(input: &str, output: &str, start: f64, duration: f64) -> anyhow::Result<()> {
140    let args = vec![
141        "-y".to_string(),
142        "-ss".into(),
143        format!("{start:.6}"),
144        "-i".into(),
145        input.into(),
146        "-t".into(),
147        format!("{duration:.6}"),
148        "-c".into(),
149        "copy".into(),
150        "-avoid_negative_ts".into(),
151        "make_zero".into(),
152        output.into(),
153    ];
154
155    let output = Command::new(ffmpeg_path())
156        .args(&args)
157        .stderr(std::process::Stdio::piped())
158        .output()
159        .await?;
160
161    if !output.status.success() {
162        let stderr = String::from_utf8_lossy(&output.stderr);
163        anyhow::bail!("ffmpeg extract failed: {stderr}");
164    }
165    Ok(())
166}
167
168/// Concatenates multiple encoded chunks into a single output without re-encoding.
169pub async fn concat(inputs: &[String], output: &str) -> anyhow::Result<()> {
170    if inputs.is_empty() {
171        anyhow::bail!("cannot concat an empty input list");
172    }
173
174    let list_path = make_concat_list_path(output);
175    let list_body = inputs
176        .iter()
177        .map(|path| format!("file '{}'", path.replace('\'', "'\\''")))
178        .collect::<Vec<_>>()
179        .join("\n");
180    std::fs::write(&list_path, format!("{list_body}\n"))?;
181
182    let args = vec![
183        "-y".to_string(),
184        "-f".into(),
185        "concat".into(),
186        "-safe".into(),
187        "0".into(),
188        "-i".into(),
189        list_path.to_string_lossy().into_owned(),
190        "-c".into(),
191        "copy".into(),
192        output.into(),
193    ];
194
195    let result = run_ffmpeg(args, None).await;
196    let _ = std::fs::remove_file(&list_path);
197    result
198}
199
200enum EncodePass<'a> {
201    Single,
202    First(&'a Path),
203    Second(&'a Path),
204}
205
206fn build_encode_args(job: &EncodeJob, pass: EncodePass<'_>) -> anyhow::Result<Vec<String>> {
207    let mut args = vec!["-y".into(), "-i".into(), job.input.clone(), "-an".into()];
208
209    if !matches!(pass, EncodePass::First(_)) {
210        args.extend(["-progress".into(), "pipe:1".into(), "-nostats".into()]);
211    }
212
213    args.extend(["-c:v".into(), job.codec.as_str().into()]);
214
215    // Rate control mode
216    match job.rate_control {
217        RateControlMode::Qp => match job.codec {
218            Codec::SvtAv1 => {
219                args.extend(["-qp".into(), job.crf.to_string()]);
220                args.extend(["-svtav1-params".into(), "enable-adaptive-quantization=0".into()]);
221            }
222            _ => {
223                args.extend(["-qp".into(), job.crf.to_string()]);
224            }
225        },
226        RateControlMode::CappedCrf => {
227            if job.max_bitrate <= 0.0 {
228                anyhow::bail!("max bitrate must be greater than zero for capped CRF mode");
229            }
230            let bufsize = if job.bufsize > 0.0 { job.bufsize } else { job.max_bitrate * 2.0 };
231            args.extend(["-crf".into(), job.crf.to_string()]);
232            args.extend(["-maxrate".into(), format!("{:.0}k", job.max_bitrate)]);
233            args.extend(["-bufsize".into(), format!("{:.0}k", bufsize)]);
234        }
235        RateControlMode::Vbr => {
236            if job.target_bitrate <= 0.0 {
237                anyhow::bail!("target bitrate must be greater than zero for VBR mode");
238            }
239            args.extend(["-b:v".into(), format!("{:.0}k", job.target_bitrate)]);
240            args.extend(["-maxrate".into(), format!("{:.0}k", job.target_bitrate * 2.0)]);
241            args.extend(["-bufsize".into(), format!("{:.0}k", job.target_bitrate * 4.0)]);
242
243            let passlog = match pass {
244                EncodePass::First(path) => {
245                    args.extend(["-pass".into(), "1".into()]);
246                    path
247                }
248                EncodePass::Second(path) => {
249                    args.extend(["-pass".into(), "2".into()]);
250                    path
251                }
252                EncodePass::Single => {
253                    anyhow::bail!("VBR mode requires a two-pass encode flow");
254                }
255            };
256
257            args.extend(["-passlogfile".into(), passlog.to_string_lossy().into_owned()]);
258        }
259        RateControlMode::Crf => {
260            args.extend(["-crf".into(), job.crf.to_string()]);
261        }
262    }
263
264    if !job.preset.is_empty() {
265        args.extend(["-preset".into(), job.preset.clone()]);
266    }
267
268    if let Some(ref res) = job.resolution {
269        if res.width > 0 && res.height > 0 {
270            args.extend([
271                "-vf".into(),
272                format!("scale={}:{}:flags=lanczos", res.width, res.height),
273            ]);
274        }
275    }
276
277    args.extend(job.extra_args.iter().cloned());
278
279    match pass {
280        EncodePass::First(_) => {
281            args.extend(["-f".into(), "null".into()]);
282            args.push(null_output_path().into());
283        }
284        EncodePass::Single | EncodePass::Second(_) => args.push(job.output.clone()),
285    }
286
287    Ok(args)
288}
289
290fn make_passlog_prefix(output: &str) -> PathBuf {
291    let output_path = Path::new(output);
292    let parent =
293        output_path.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or(Path::new("."));
294    let stem = output_path.file_stem().and_then(|s| s.to_str()).unwrap_or("viser");
295    let unique = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis()).unwrap_or(0);
296    parent.join(format!(".{stem}.viser-passlog-{unique}-{}", std::process::id()))
297}
298
299fn make_concat_list_path(output: &str) -> PathBuf {
300    let output_path = Path::new(output);
301    let parent =
302        output_path.parent().filter(|p| !p.as_os_str().is_empty()).unwrap_or(Path::new("."));
303    let stem = output_path.file_stem().and_then(|s| s.to_str()).unwrap_or("viser");
304    let unique = SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_millis()).unwrap_or(0);
305    parent.join(format!(".{stem}.viser-concat-{unique}-{}.txt", std::process::id()))
306}
307
308fn null_output_path() -> &'static str {
309    if cfg!(windows) { "NUL" } else { "/dev/null" }
310}
311
312struct PasslogCleanup {
313    parent: PathBuf,
314    prefix: String,
315}
316
317impl PasslogCleanup {
318    fn new(path: PathBuf) -> Self {
319        let parent = path.parent().unwrap_or(Path::new(".")).to_path_buf();
320        let prefix = path.file_name().and_then(|name| name.to_str()).unwrap_or_default().to_owned();
321        Self { parent, prefix }
322    }
323
324    fn run(&self) {
325        let Ok(entries) = std::fs::read_dir(&self.parent) else {
326            return;
327        };
328
329        for entry in entries.flatten() {
330            let path = entry.path();
331            let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
332                continue;
333            };
334            if !name.starts_with(&self.prefix) {
335                continue;
336            }
337            if let Err(err) = std::fs::remove_file(&path) {
338                tracing::debug!(?path, ?err, "failed to remove ffmpeg two-pass log file");
339            }
340        }
341    }
342}
343
344/// Returns true when a complete progress block is ready.
345fn parse_progress_line(line: &str, p: &mut Progress) -> bool {
346    let Some((key, value)) = line.split_once('=') else {
347        return false;
348    };
349
350    match key {
351        "frame" => {
352            p.frame = value.parse().unwrap_or(0);
353        }
354        "fps" => {
355            p.fps = value.parse().unwrap_or(0.0);
356        }
357        "bitrate" => {
358            let v = value.trim_end_matches("kbits/s");
359            p.bitrate = v.parse().unwrap_or(0.0);
360        }
361        "speed" => {
362            let v = value.trim_end_matches('x');
363            p.speed = v.parse().unwrap_or(0.0);
364        }
365        "out_time_us" => {
366            let us: u64 = value.parse().unwrap_or(0);
367            p.time = Duration::from_micros(us);
368        }
369        "progress" => return true,
370        _ => {}
371    }
372    false
373}
374
375#[cfg(test)]
376mod tests {
377    use super::*;
378    use crate::Codec;
379
380    fn sample_job(mode: RateControlMode) -> EncodeJob {
381        EncodeJob {
382            input: "input.mp4".into(),
383            output: "output.mp4".into(),
384            resolution: Some(crate::Resolution::new(1280, 720)),
385            codec: Codec::X264,
386            crf: 23,
387            rate_control: mode,
388            target_bitrate: 2500.0,
389            max_bitrate: 3000.0,
390            bufsize: 6000.0,
391            preset: "medium".into(),
392            extra_args: vec![],
393        }
394    }
395
396    #[test]
397    fn test_build_encode_args_crf_single_pass() {
398        let args =
399            build_encode_args(&sample_job(RateControlMode::Crf), EncodePass::Single).unwrap();
400        assert!(args.windows(2).any(|w| w == ["-crf", "23"]));
401        assert_eq!(args.last().unwrap(), "output.mp4");
402    }
403
404    #[test]
405    fn test_build_encode_args_vbr_first_pass_uses_null_output() {
406        let job = sample_job(RateControlMode::Vbr);
407        let passlog = Path::new("/tmp/viser-passlog");
408        let args = build_encode_args(&job, EncodePass::First(passlog)).unwrap();
409        assert!(args.windows(2).any(|w| w == ["-pass", "1"]));
410        assert!(args.windows(2).any(|w| w == ["-f", "null"]));
411        assert_eq!(args.last().unwrap(), null_output_path());
412    }
413
414    #[test]
415    fn test_build_encode_args_vbr_second_pass_writes_output() {
416        let job = sample_job(RateControlMode::Vbr);
417        let passlog = Path::new("/tmp/viser-passlog");
418        let args = build_encode_args(&job, EncodePass::Second(passlog)).unwrap();
419        assert!(args.windows(2).any(|w| w == ["-pass", "2"]));
420        assert_eq!(args.last().unwrap(), "output.mp4");
421    }
422
423    #[test]
424    fn test_build_encode_args_capped_crf_sets_vbv() {
425        let args =
426            build_encode_args(&sample_job(RateControlMode::CappedCrf), EncodePass::Single).unwrap();
427        assert!(args.windows(2).any(|w| w == ["-crf", "23"]));
428        assert!(args.windows(2).any(|w| w == ["-maxrate", "3000k"]));
429        assert!(args.windows(2).any(|w| w == ["-bufsize", "6000k"]));
430    }
431}