1use crate::sink::FrameSink;
18use crate::wl::CapturedImage;
19use anyhow::{Context, Result, anyhow, bail};
20use ffmpeg::format::Pixel;
21use ffmpeg_next as ffmpeg;
22use std::path::{Path, PathBuf};
23use std::sync::Once;
24use std::time::Duration;
25
26static FFMPEG_INIT: Once = Once::new();
27
28fn ensure_ffmpeg() {
31 FFMPEG_INIT.call_once(|| {
32 let _ = ffmpeg::init();
34 ffmpeg::util::log::set_level(ffmpeg::util::log::Level::Warning);
35 });
36}
37
38#[derive(Clone, Copy, Debug, PartialEq, Eq)]
41pub enum Backend {
42 Auto,
44 Nvenc,
46 Vaapi,
48 Software,
50}
51
52impl Backend {
53 fn codec_name(self) -> &'static str {
55 match self {
56 Backend::Nvenc => "h264_nvenc",
57 Backend::Vaapi => "h264_vaapi",
58 Backend::Software => "libx264",
59 Backend::Auto => unreachable!("resolved before use"),
60 }
61 }
62}
63
64#[derive(Clone, Copy, Debug, PartialEq, Eq)]
66pub enum Mode {
67 Record,
69 Timelapse,
72}
73
74#[derive(Clone, Debug)]
76pub struct Options {
77 pub backend: Backend,
79 pub fps: u32,
81 pub mode: Mode,
83 pub device: Option<PathBuf>,
85 pub audio: bool,
88}
89
90impl Default for Options {
91 fn default() -> Self {
92 Self {
93 backend: Backend::Auto,
94 fps: 30,
95 mode: Mode::Record,
96 device: None,
97 audio: false,
98 }
99 }
100}
101
102const MS_TIMEBASE: ffmpeg::Rational = ffmpeg::Rational(1, 1000);
104pub(crate) const AUDIO_RATE: u32 = 48_000;
106pub(crate) const AUDIO_CHANNELS: usize = 2;
107const AUDIO_BIT_RATE: usize = 160_000;
109
110struct Pipeline {
112 octx: ffmpeg::format::context::Output,
113 encoder: ffmpeg::encoder::Video,
114 scaler: ffmpeg::software::scaling::Context,
115 src: (u32, u32),
117 dst: (u32, u32),
119 enc_time_base: ffmpeg::Rational,
120 ost_time_base: ffmpeg::Rational,
121 target_format: Pixel,
122 last_pts: i64,
124 index: i64,
126 vaapi: Option<VaapiCtx>,
128 audio: Option<AudioPipe>,
130}
131
132struct AudioPipe {
135 encoder: ffmpeg::encoder::Audio,
136 stream_index: usize,
137 enc_time_base: ffmpeg::Rational,
138 ost_time_base: ffmpeg::Rational,
139 frame_size: usize,
141 pts: i64,
143}
144
145struct VaapiCtx {
148 device: *mut ffmpeg::ffi::AVBufferRef,
149 frames: *mut ffmpeg::ffi::AVBufferRef,
150}
151
152impl VaapiCtx {
153 fn new(device: Option<&Path>, w: u32, h: u32) -> Result<Self> {
157 use ffmpeg::ffi;
158 use std::os::unix::ffi::OsStrExt;
159
160 let cpath = match device {
161 Some(p) => Some(
162 std::ffi::CString::new(p.as_os_str().as_bytes())
163 .context("device path contains a NUL byte")?,
164 ),
165 None => None,
166 };
167 let dptr = cpath.as_ref().map_or(std::ptr::null(), |c| c.as_ptr());
168
169 unsafe {
170 let mut dev: *mut ffi::AVBufferRef = std::ptr::null_mut();
171 let r = ffi::av_hwdevice_ctx_create(
172 &mut dev,
173 ffi::AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI,
174 dptr,
175 std::ptr::null_mut(),
176 0,
177 );
178 if r < 0 {
179 let name =
180 device.map_or_else(|| "(default)".to_string(), |p| p.display().to_string());
181 bail!("opening VAAPI device {name} (code {r})");
182 }
183
184 let frames = ffi::av_hwframe_ctx_alloc(dev);
185 if frames.is_null() {
186 ffi::av_buffer_unref(&mut dev);
187 bail!("allocating the VAAPI frame pool");
188 }
189 let fctx = (*frames).data as *mut ffi::AVHWFramesContext;
190 (*fctx).format = ffi::AVPixelFormat::AV_PIX_FMT_VAAPI;
191 (*fctx).sw_format = ffi::AVPixelFormat::AV_PIX_FMT_NV12;
192 (*fctx).width = w as i32;
193 (*fctx).height = h as i32;
194 (*fctx).initial_pool_size = 20;
195
196 let r = ffi::av_hwframe_ctx_init(frames);
197 if r < 0 {
198 let mut frames = frames;
199 ffi::av_buffer_unref(&mut frames);
200 ffi::av_buffer_unref(&mut dev);
201 bail!("initialising the VAAPI frame pool (code {r})");
202 }
203 Ok(Self {
204 device: dev,
205 frames,
206 })
207 }
208 }
209}
210
211impl Drop for VaapiCtx {
212 fn drop(&mut self) {
213 unsafe {
215 ffmpeg::ffi::av_buffer_unref(&mut self.frames);
216 ffmpeg::ffi::av_buffer_unref(&mut self.device);
217 }
218 }
219}
220
221pub struct VideoEncoder {
223 path: PathBuf,
224 opts: Options,
225 pipeline: Option<Pipeline>,
226 audio_buf: Vec<f32>,
229}
230
231impl VideoEncoder {
232 pub fn new(path: impl Into<PathBuf>, opts: Options) -> Result<Self> {
235 ensure_ffmpeg();
236 Ok(Self {
237 path: path.into(),
238 opts,
239 pipeline: None,
240 audio_buf: Vec::new(),
241 })
242 }
243
244 pub fn resolved_backend(&self) -> Result<Backend> {
246 resolve_backend(self.opts.backend)
247 }
248}
249
250fn resolve_backend(backend: Backend) -> Result<Backend> {
252 ensure_ffmpeg();
253 let available = |b: Backend| ffmpeg::encoder::find_by_name(b.codec_name()).is_some();
254 match backend {
255 Backend::Auto => [Backend::Nvenc, Backend::Vaapi, Backend::Software]
256 .into_iter()
257 .find(|&b| available(b))
258 .ok_or_else(|| anyhow!("no H.264 encoder available (need NVENC, VAAPI or libx264)")),
259 b if available(b) => Ok(b),
260 b => bail!(
261 "encoder '{}' is not available in this FFmpeg build",
262 b.codec_name()
263 ),
264 }
265}
266
267fn build_audio_stream(
269 octx: &mut ffmpeg::format::context::Output,
270 global_header: bool,
271) -> Result<AudioPipe> {
272 let codec = ffmpeg::encoder::find(ffmpeg::codec::Id::AAC)
273 .ok_or_else(|| anyhow!("no AAC encoder in this FFmpeg build"))?;
274 let mut astream = octx.add_stream(codec).context("adding audio stream")?;
275 let stream_index = astream.index();
276
277 let mut aenc = ffmpeg::codec::context::Context::new_with_codec(codec)
278 .encoder()
279 .audio()?;
280 aenc.set_rate(AUDIO_RATE as i32);
281 aenc.set_channel_layout(ffmpeg::channel_layout::ChannelLayout::STEREO);
282 aenc.set_format(ffmpeg::format::Sample::F32(
283 ffmpeg::format::sample::Type::Planar,
284 ));
285 aenc.set_bit_rate(AUDIO_BIT_RATE);
286 let enc_time_base = ffmpeg::Rational(1, AUDIO_RATE as i32);
287 aenc.set_time_base(enc_time_base);
288 if global_header {
289 aenc.set_flags(ffmpeg::codec::Flags::GLOBAL_HEADER);
290 }
291
292 let encoder = aenc.open_as(codec).context("opening the AAC encoder")?;
293 astream.set_parameters(&encoder);
294 let frame_size = (encoder.frame_size() as usize).max(1);
295
296 Ok(AudioPipe {
297 encoder,
298 stream_index,
299 enc_time_base,
300 ost_time_base: enc_time_base, frame_size,
302 pts: 0,
303 })
304}
305
306impl Pipeline {
307 fn new(path: &Path, opts: &Options, sw: u32, sh: u32) -> Result<Self> {
309 let backend = resolve_backend(opts.backend)?;
310 let codec = ffmpeg::encoder::find_by_name(backend.codec_name())
311 .ok_or_else(|| anyhow!("encoder '{}' unavailable", backend.codec_name()))?;
312
313 let dst = (sw & !1, sh & !1);
315 if dst.0 == 0 || dst.1 == 0 {
316 bail!("source too small to encode ({sw}x{sh})");
317 }
318 let (enc_format, target_format) = match backend {
322 Backend::Software => (Pixel::YUV420P, Pixel::YUV420P),
323 Backend::Nvenc => (Pixel::NV12, Pixel::NV12),
324 Backend::Vaapi => (Pixel::VAAPI, Pixel::NV12),
325 Backend::Auto => unreachable!("resolved above"),
326 };
327
328 let mut octx = ffmpeg::format::output(&path)
329 .with_context(|| format!("opening output '{}'", path.display()))?;
330 let global_header = octx
331 .format()
332 .flags()
333 .contains(ffmpeg::format::Flags::GLOBAL_HEADER);
334
335 let mut ost = octx.add_stream(codec).context("adding video stream")?;
336 let mut enc = ffmpeg::codec::context::Context::new_with_codec(codec)
337 .encoder()
338 .video()?;
339 enc.set_width(dst.0);
340 enc.set_height(dst.1);
341 enc.set_format(enc_format);
342 enc.set_frame_rate(Some(ffmpeg::Rational(opts.fps as i32, 1)));
343 let enc_time_base = match opts.mode {
345 Mode::Record => MS_TIMEBASE,
346 Mode::Timelapse => ffmpeg::Rational(1, opts.fps as i32),
347 };
348 enc.set_time_base(enc_time_base);
349 if global_header {
350 enc.set_flags(ffmpeg::codec::Flags::GLOBAL_HEADER);
351 }
352
353 let vaapi = if backend == Backend::Vaapi {
355 let ctx =
356 VaapiCtx::new(opts.device.as_deref(), dst.0, dst.1).context("setting up VAAPI")?;
357 unsafe {
358 (*enc.as_mut_ptr()).hw_frames_ctx = ffmpeg::ffi::av_buffer_ref(ctx.frames);
359 }
360 Some(ctx)
361 } else {
362 None
363 };
364
365 let encoder = enc
366 .open_as(codec)
367 .with_context(|| format!("opening encoder '{}'", backend.codec_name()))?;
368 ost.set_parameters(&encoder);
369
370 let mut audio = if opts.audio && opts.mode == Mode::Record {
373 Some(build_audio_stream(&mut octx, global_header)?)
374 } else {
375 None
376 };
377
378 octx.write_header().context("writing container header")?;
379 let ost_time_base = octx.stream(0).context("no output stream")?.time_base();
381 if let Some(ap) = audio.as_mut() {
382 ap.ost_time_base = octx
383 .stream(ap.stream_index)
384 .context("no audio stream")?
385 .time_base();
386 }
387
388 let scaler = ffmpeg::software::scaling::Context::get(
389 Pixel::RGBA,
390 sw,
391 sh,
392 target_format,
393 dst.0,
394 dst.1,
395 ffmpeg::software::scaling::Flags::BILINEAR,
396 )
397 .context("creating RGBA->YUV scaler")?;
398
399 Ok(Self {
400 octx,
401 encoder,
402 scaler,
403 src: (sw, sh),
404 dst,
405 enc_time_base,
406 ost_time_base,
407 target_format,
408 last_pts: -1,
409 index: 0,
410 vaapi,
411 audio,
412 })
413 }
414
415 fn ensure_scaler(&mut self, sw: u32, sh: u32) -> Result<()> {
417 if self.src == (sw, sh) {
418 return Ok(());
419 }
420 self.scaler = ffmpeg::software::scaling::Context::get(
421 Pixel::RGBA,
422 sw,
423 sh,
424 self.target_format,
425 self.dst.0,
426 self.dst.1,
427 ffmpeg::software::scaling::Flags::BILINEAR,
428 )
429 .context("rebuilding scaler for new source size")?;
430 self.src = (sw, sh);
431 Ok(())
432 }
433
434 fn encode(&mut self, img: &CapturedImage, ts: Duration, mode: Mode) -> Result<()> {
436 if img.width == 0 || img.height == 0 {
437 return Ok(());
438 }
439 self.ensure_scaler(img.width, img.height)?;
440
441 let mut src = ffmpeg::frame::Video::new(Pixel::RGBA, img.width, img.height);
442 copy_rgba_into(&mut src, img);
443
444 let mut dst = ffmpeg::frame::Video::new(self.target_format, self.dst.0, self.dst.1);
445 self.scaler.run(&src, &mut dst).context("scaling frame")?;
446
447 let pts = match mode {
448 Mode::Record => (ts.as_millis() as i64).max(self.last_pts + 1),
449 Mode::Timelapse => self.index,
450 };
451 self.last_pts = pts;
452 self.index += 1;
453
454 if let Some(vaapi) = &self.vaapi {
455 let mut hw = ffmpeg::frame::Video::empty();
457 unsafe {
458 let r = ffmpeg::ffi::av_hwframe_get_buffer(vaapi.frames, hw.as_mut_ptr(), 0);
459 if r < 0 {
460 bail!("allocating a VAAPI surface (code {r})");
461 }
462 let r = ffmpeg::ffi::av_hwframe_transfer_data(hw.as_mut_ptr(), dst.as_ptr(), 0);
463 if r < 0 {
464 bail!("uploading the frame to the GPU (code {r})");
465 }
466 }
467 hw.set_pts(Some(pts));
468 self.encoder.send_frame(&hw).context("sending frame")?;
469 } else {
470 dst.set_pts(Some(pts));
471 self.encoder.send_frame(&dst).context("sending frame")?;
472 }
473 self.drain()
474 }
475
476 fn drain(&mut self) -> Result<()> {
478 let mut packet = ffmpeg::Packet::empty();
479 while self.encoder.receive_packet(&mut packet).is_ok() {
480 packet.set_stream(0);
481 packet.rescale_ts(self.enc_time_base, self.ost_time_base);
482 packet
483 .write_interleaved(&mut self.octx)
484 .context("writing packet")?;
485 }
486 Ok(())
487 }
488
489 fn encode_audio(&mut self, buf: &mut Vec<f32>) -> Result<()> {
492 let frame_size = match &self.audio {
493 Some(a) => a.frame_size,
494 None => {
495 buf.clear();
496 return Ok(());
497 }
498 };
499 let need = frame_size * AUDIO_CHANNELS;
500 while buf.len() >= need {
501 let mut planes: Vec<Vec<f32>> = (0..AUDIO_CHANNELS)
502 .map(|_| Vec::with_capacity(frame_size))
503 .collect();
504 for fr in buf[..need].chunks_exact(AUDIO_CHANNELS) {
505 for (c, p) in planes.iter_mut().enumerate() {
506 p.push(fr[c]);
507 }
508 }
509 buf.drain(..need);
510
511 let mut frame = ffmpeg::frame::Audio::new(
512 ffmpeg::format::Sample::F32(ffmpeg::format::sample::Type::Planar),
513 frame_size,
514 ffmpeg::channel_layout::ChannelLayout::STEREO,
515 );
516 frame.set_rate(AUDIO_RATE);
517 for (c, p) in planes.iter().enumerate() {
518 frame.plane_mut::<f32>(c).copy_from_slice(p);
519 }
520
521 let ap = self.audio.as_mut().expect("audio present");
522 frame.set_pts(Some(ap.pts));
523 ap.pts += frame_size as i64;
524 ap.encoder
525 .send_frame(&frame)
526 .context("sending audio frame")?;
527
528 let mut packet = ffmpeg::Packet::empty();
529 while ap.encoder.receive_packet(&mut packet).is_ok() {
530 packet.set_stream(ap.stream_index);
531 packet.rescale_ts(ap.enc_time_base, ap.ost_time_base);
532 packet
533 .write_interleaved(&mut self.octx)
534 .context("writing audio packet")?;
535 }
536 }
537 Ok(())
538 }
539
540 fn finish(mut self) -> Result<()> {
542 self.encoder.send_eof().context("flushing encoder")?;
543 self.drain()?;
544 if let Some(ap) = self.audio.as_mut() {
545 ap.encoder.send_eof().context("flushing audio encoder")?;
546 let mut packet = ffmpeg::Packet::empty();
547 while ap.encoder.receive_packet(&mut packet).is_ok() {
548 packet.set_stream(ap.stream_index);
549 packet.rescale_ts(ap.enc_time_base, ap.ost_time_base);
550 packet
551 .write_interleaved(&mut self.octx)
552 .context("writing final audio packet")?;
553 }
554 }
555 self.octx
556 .write_trailer()
557 .context("writing container trailer")?;
558 Ok(())
559 }
560}
561
562fn copy_rgba_into(frame: &mut ffmpeg::frame::Video, img: &CapturedImage) {
565 let w = img.width as usize;
566 let stride = frame.stride(0);
567 let row_bytes = w * 4;
568 let dst = frame.data_mut(0);
569 for y in 0..img.height as usize {
570 let s = y * row_bytes;
571 let d = y * stride;
572 dst[d..d + row_bytes].copy_from_slice(&img.rgba[s..s + row_bytes]);
573 }
574}
575
576impl FrameSink for VideoEncoder {
577 fn push(&mut self, img: &CapturedImage, ts: Duration) -> Result<()> {
578 if self.pipeline.is_none() {
579 self.pipeline = Some(Pipeline::new(
580 &self.path, &self.opts, img.width, img.height,
581 )?);
582 }
583 let mode = self.opts.mode;
584 let p = self.pipeline.as_mut().expect("just initialised");
585 p.encode(img, ts, mode)?;
586 p.encode_audio(&mut self.audio_buf)
587 }
588
589 fn push_audio(&mut self, pcm: &[f32]) {
593 if !self.opts.audio {
594 return;
595 }
596 self.audio_buf.extend_from_slice(pcm);
597 let cap = AUDIO_RATE as usize * AUDIO_CHANNELS * 5;
598 if self.audio_buf.len() > cap {
599 let drop = self.audio_buf.len() - cap;
600 self.audio_buf.drain(..drop);
601 }
602 }
603
604 fn finish(&mut self) -> Result<()> {
605 match self.pipeline.take() {
606 Some(mut p) => {
607 p.encode_audio(&mut self.audio_buf)?; p.finish()
609 }
610 None => Ok(()), }
612 }
613}
614
615#[cfg(test)]
616mod tests {
617 use super::*;
618
619 fn frame(w: u32, h: u32, t: u32) -> CapturedImage {
622 let mut rgba = vec![0u8; (w * h * 4) as usize];
623 for y in 0..h {
624 for x in 0..w {
625 let i = ((y * w + x) * 4) as usize;
626 rgba[i] = ((x + t) & 0xff) as u8;
627 rgba[i + 1] = ((y + t) & 0xff) as u8;
628 rgba[i + 2] = ((x + y) & 0xff) as u8;
629 rgba[i + 3] = 255;
630 }
631 }
632 CapturedImage {
633 width: w,
634 height: h,
635 rgba,
636 }
637 }
638
639 fn run_encode(requested: Backend) {
643 let backend = match resolve_backend(requested) {
644 Ok(b) => b,
645 Err(_) => {
646 eprintln!("backend {requested:?} unavailable; skipping");
647 return;
648 }
649 };
650
651 let (w, h, fps, n) = (320u32, 240u32, 30u32, 30u32);
652 let path = std::env::temp_dir().join(format!(
654 "wlr_capture_enc_{}_{}.mp4",
655 std::process::id(),
656 backend.codec_name()
657 ));
658 let mut enc = VideoEncoder::new(
659 &path,
660 Options {
661 backend,
662 fps,
663 mode: Mode::Record,
664 device: Some("/dev/dri/renderD128".into()),
665 audio: false,
666 },
667 )
668 .expect("create encoder");
669
670 for i in 0..n {
671 let ts = Duration::from_millis((i * 1000 / fps) as u64);
672 enc.push(&frame(w, h, i), ts).expect("push frame");
673 }
674 enc.finish().expect("finish");
675
676 let meta = std::fs::metadata(&path).expect("output file exists");
677 assert!(
678 meta.len() > 1000,
679 "output suspiciously small: {} bytes",
680 meta.len()
681 );
682
683 if let Ok(out) = std::process::Command::new("ffprobe")
685 .args([
686 "-v",
687 "error",
688 "-select_streams",
689 "v:0",
690 "-show_entries",
691 "stream=codec_name,width,height",
692 "-of",
693 "default=nw=1:nk=1",
694 ])
695 .arg(&path)
696 .output()
697 && out.status.success()
698 {
699 let s = String::from_utf8_lossy(&out.stdout);
700 let fields: Vec<&str> = s.split_whitespace().collect();
701 assert_eq!(fields, ["h264", "320", "240"], "ffprobe stream metadata");
702 }
703
704 let _ = std::fs::remove_file(&path);
705 }
706
707 #[test]
709 fn encodes_software() {
710 run_encode(Backend::Software);
711 }
712
713 #[test]
720 #[ignore]
721 fn encodes_nvenc() {
722 run_encode(Backend::Nvenc);
723 }
724
725 #[test]
728 #[ignore]
729 fn encodes_vaapi() {
730 run_encode(Backend::Vaapi);
731 }
732}