wlr_capture/
sink.rs

1//! Common output seam for capture-consuming tools (screenshot, record, timelapse).
2//!
3//! A capture round yields a [`Frame`] — CPU pixels (shm) or a GPU dma-buf
4//! (zero-copy). A [`FrameSink`] consumes a stream of those frames. Most sinks only
5//! want CPU pixels, so the default [`FrameSink::push_dmabuf`] reads the dma-buf back
6//! via a [`GpuReadback`] and forwards to [`FrameSink::push`]; a GPU-native sink (a
7//! future hardware video encoder, say) can override it to keep the buffer on the
8//! GPU. [`pump`] routes a `Frame` to whichever path applies, creating the readback
9//! lazily so pure-shm streams never spin up an EGL context.
10
11use crate::gl::GpuReadback;
12use crate::wl::{CapturedImage, DmabufFrame, Frame};
13use anyhow::Result;
14use std::time::Duration;
15
16/// A consumer of a capture stream. `ts` is each frame's capture time relative to a
17/// start the sink defines (a screenshot ignores it; a recorder uses it for timing).
18pub trait FrameSink {
19    /// Consume one CPU-pixel frame.
20    fn push(&mut self, img: &CapturedImage, ts: Duration) -> Result<()>;
21
22    /// Feed interleaved PCM for an optional audio track. The default ignores it (sinks
23    /// without sound — GIF/WebP — and silent recordings); the video encoder buffers it
24    /// and muxes it on the next [`FrameSink::push`].
25    fn push_audio(&mut self, _pcm: &[f32]) {}
26
27    /// Consume one GPU dma-buf frame. The default reads it back to CPU pixels via
28    /// `rb` and forwards to [`FrameSink::push`]; override to consume it on the GPU.
29    fn push_dmabuf(
30        &mut self,
31        rb: &mut GpuReadback,
32        frame: DmabufFrame,
33        ts: Duration,
34    ) -> Result<()> {
35        let img = rb.readback(frame)?;
36        self.push(&img, ts)
37    }
38
39    /// Flush and finalize (write the file, close the encoder, …). Call once, last.
40    fn finish(&mut self) -> Result<()> {
41        Ok(())
42    }
43}
44
45/// Route one [`Frame`] to `sink`, picking the CPU or dma-buf path. The readback
46/// context lives in `rb` and is built on first need — a stream that only ever
47/// produces shm frames (a no-GPU build) never constructs one. Hold a single
48/// `Option<GpuReadback>` across the whole stream so the context is reused.
49pub fn pump(
50    sink: &mut dyn FrameSink,
51    rb: &mut Option<GpuReadback>,
52    frame: Frame,
53    ts: Duration,
54) -> Result<()> {
55    match frame {
56        Frame::Shm(img) => sink.push(&img, ts),
57        Frame::Dmabuf(d) => {
58            let rb = match rb {
59                Some(rb) => rb,
60                None => rb.insert(GpuReadback::new()?),
61            };
62            sink.push_dmabuf(rb, d, ts)
63        }
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    /// A sink that records what it received — exercises the shm path of [`pump`]
72    /// and the trait's stamping without needing a GPU.
73    #[derive(Default)]
74    struct Collect {
75        frames: Vec<(u32, u32, Duration)>,
76        finished: bool,
77    }
78
79    impl FrameSink for Collect {
80        fn push(&mut self, img: &CapturedImage, ts: Duration) -> Result<()> {
81            self.frames.push((img.width, img.height, ts));
82            Ok(())
83        }
84        fn finish(&mut self) -> Result<()> {
85            self.finished = true;
86            Ok(())
87        }
88    }
89
90    #[test]
91    fn pump_routes_shm_frames() {
92        let mut sink = Collect::default();
93        let mut rb = None; // never built: no dma-buf frames in this stream
94        let img = CapturedImage {
95            width: 4,
96            height: 2,
97            rgba: vec![0; 4 * 2 * 4],
98        };
99        pump(
100            &mut sink,
101            &mut rb,
102            Frame::Shm(img),
103            Duration::from_millis(40),
104        )
105        .unwrap();
106        sink.finish().unwrap();
107
108        assert!(rb.is_none());
109        assert_eq!(sink.frames, vec![(4, 2, Duration::from_millis(40))]);
110        assert!(sink.finished);
111    }
112}
wlr_capture/sink.rs

wlr_capture/
sink.rs
