Skip to main content

ustreamer_capture/
staging.rs

1//! Triple-buffered staging buffer capture (cross-platform fallback).
2
3use crate::{CaptureError, CapturedFrame, FrameCapture};
4
5/// Cross-platform frame capture using GPU→CPU staging buffers.
6///
7/// Uses a ring of staging buffers to overlap capture of frame N
8/// with encoding of frame N-1 and rendering of frame N+1.
9pub struct StagingCapture {
10    buffers: Vec<Option<wgpu::Buffer>>,
11    current: usize,
12    capacity: usize,
13}
14
15impl StagingCapture {
16    /// Create a new staging capture with the given number of ring buffers.
17    pub fn new(ring_size: usize) -> Self {
18        Self {
19            buffers: (0..ring_size).map(|_| None).collect(),
20            current: 0,
21            capacity: ring_size,
22        }
23    }
24}
25
26impl FrameCapture for StagingCapture {
27    fn capture(
28        &mut self,
29        _instance: &wgpu::Instance,
30        device: &wgpu::Device,
31        queue: &wgpu::Queue,
32        texture: &wgpu::Texture,
33    ) -> Result<CapturedFrame, CaptureError> {
34        let size = texture.size();
35        let format = texture.format();
36
37        let bytes_per_pixel = format
38            .block_copy_size(None)
39            .ok_or(CaptureError::UnsupportedFormat(format))?;
40
41        // wgpu requires rows to be aligned to 256 bytes
42        let unpadded_row_bytes = size.width * bytes_per_pixel;
43        let padded_row_bytes = (unpadded_row_bytes + 255) & !255;
44        let buffer_size = (padded_row_bytes * size.height) as u64;
45
46        // Reuse or create buffer
47        let buffer = self.buffers[self.current].get_or_insert_with(|| {
48            device.create_buffer(&wgpu::BufferDescriptor {
49                label: Some("ustreamer-staging"),
50                size: buffer_size,
51                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
52                mapped_at_creation: false,
53            })
54        });
55
56        // Copy texture to staging buffer
57        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
58            label: Some("ustreamer-capture"),
59        });
60
61        encoder.copy_texture_to_buffer(
62            wgpu::TexelCopyTextureInfo {
63                texture,
64                mip_level: 0,
65                origin: wgpu::Origin3d::ZERO,
66                aspect: wgpu::TextureAspect::All,
67            },
68            wgpu::TexelCopyBufferInfo {
69                buffer,
70                layout: wgpu::TexelCopyBufferLayout {
71                    offset: 0,
72                    bytes_per_row: Some(padded_row_bytes),
73                    rows_per_image: Some(size.height),
74                },
75            },
76            texture.size(),
77        );
78
79        queue.submit(std::iter::once(encoder.finish()));
80
81        // Map and read
82        let slice = buffer.slice(..);
83        let (tx, rx) = std::sync::mpsc::channel();
84        slice.map_async(wgpu::MapMode::Read, move |result| {
85            tx.send(result).ok();
86        });
87        device.poll(wgpu::PollType::wait_indefinitely()).ok();
88
89        rx.recv()
90            .map_err(|e| CaptureError::MapFailed(e.to_string()))?
91            .map_err(|e| CaptureError::MapFailed(e.to_string()))?;
92
93        // Copy out the data (removing row padding)
94        let mapped = slice.get_mapped_range();
95        let mut data = Vec::with_capacity((unpadded_row_bytes * size.height) as usize);
96        for row in 0..size.height {
97            let start = (row * padded_row_bytes) as usize;
98            let end = start + unpadded_row_bytes as usize;
99            data.extend_from_slice(&mapped[start..end]);
100        }
101        drop(mapped);
102        buffer.unmap();
103
104        self.current = (self.current + 1) % self.capacity;
105
106        Ok(CapturedFrame::CpuBuffer {
107            data,
108            width: size.width,
109            height: size.height,
110            stride: unpadded_row_bytes,
111            format,
112        })
113    }
114}