ustreamer_capture/
staging.rs1use crate::{CaptureError, CapturedFrame, FrameCapture};
4
5pub struct StagingCapture {
10 buffers: Vec<Option<wgpu::Buffer>>,
11 current: usize,
12 capacity: usize,
13}
14
15impl StagingCapture {
16 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 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 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 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 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 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}