feat: headless capture — public headless Window, pipelined pixel readback, fast read_pixels

Off-screen video capture improvements, in three parts:

- read_pixels 18x faster: convert BGRA->RGB from cached memory (mapped
  readback memory is uncached, ~10 MB/s scalar reads), reuse the staging
  buffer across calls, and wait on the copy's submission index instead of
  polling the whole device.

- Window::new_headless_with_setup(width, height, setup): a full-featured
  Window backed by no OS window and no swapchain (same headless canvas as
  OffscreenSurface, but exposing the whole Window API — custom renderers,
  ray tracer, snap* readbacks). Never throttled by display vsync; works
  without a display server.

- snap_begin()/snap_finish(): pipelined readback that enqueues the
  texture->buffer copy and map without waiting, so frame N's copy is
  collected after frame N+1 renders instead of stalling the GPU pipeline
  each frame. read_pixels() is refactored onto the same begin/finish path
  with unchanged behavior.

Motivation: a windowed capture loop (render_frame + blocking snap) is
vsync-locked — the blocking readback defeats swapchain frame pipelining,
so every iteration eats ~2 vblanks (~30 fps at 60 Hz) regardless of GPU
load. Headless + pipelined readback took a downstream nexus3d capture
loop from 33 to 92 fps at 640x480 (2.7 fps before the read_pixels fix).

Supersedes dimforge#397 (included here).

Author: haixuanTao

src/context/context.rs

@@ -278,6 +278,16 @@ impl Context {
         self.queue.submit(command_buffers);
     }
 
+    /// Submits command buffers to the GPU queue, returning the submission
+    /// index so callers can wait for exactly this submission (rather than the
+    /// whole device).
+    pub fn submit_indexed<I: IntoIterator<Item = wgpu::CommandBuffer>>(
+        &self,
+        command_buffers: I,
+    ) -> wgpu::SubmissionIndex {
+        self.queue.submit(command_buffers)
+    }
+
     /// Writes texture data to the GPU.
     ///
     /// # Arguments

src/window/canvas.rs

@@ -232,4 +232,16 @@ impl Canvas {
     pub fn read_pixels(&self, out: &mut Vec<u8>, x: usize, y: usize, width: usize, height: usize) {
         self.canvas.read_pixels(out, x, y, width, height)
     }
+
+    /// Starts a non-blocking readback of the readback texture; complete it
+    /// with [`Self::finish_read_pixels`]. See `WgpuCanvas::begin_read_pixels`.
+    pub fn begin_read_pixels(&self, x: usize, y: usize, width: usize, height: usize) {
+        self.canvas.begin_read_pixels(x, y, width, height)
+    }
+
+    /// Completes a readback started by [`Self::begin_read_pixels`], returning
+    /// the captured `(width, height)`, or `None` when none is in flight.
+    pub fn finish_read_pixels(&self, out: &mut Vec<u8>) -> Option<(u32, u32)> {
+        self.canvas.finish_read_pixels(out)
+    }
 }

src/window/screenshot.rs

@@ -72,4 +72,29 @@ impl Window {
         let img = img_opt.expect("Buffer created from window was not big enough for image.");
         imageops::flip_vertical(&img)
     }
+
+    /// Starts a non-blocking capture of the last rendered frame.
+    ///
+    /// This enqueues the GPU→CPU copy of the framebuffer but does not wait for
+    /// it; collect the pixels with [`Self::snap_finish`], typically after
+    /// rendering the *next* frame, so the copy overlaps with useful GPU work
+    /// instead of stalling the pipeline the way the blocking [`Self::snap`]
+    /// does. One capture can be in flight at a time; a second `snap_begin`
+    /// completes and discards the previous one.
+    pub fn snap_begin(&self) {
+        let (width, height) = self.canvas.size();
+        self.canvas
+            .begin_read_pixels(0, 0, width as usize, height as usize);
+    }
+
+    /// Completes a capture started by [`Self::snap_begin`], returning the
+    /// frame as an image (top-left origin, like [`Self::snap_image`]), or
+    /// `None` when no capture is in flight.
+    pub fn snap_finish(&self) -> Option<ImageBuffer<Rgb<u8>, Vec<u8>>> {
+        let mut buf = Vec::new();
+        let (width, height) = self.canvas.finish_read_pixels(&mut buf)?;
+        let img = ImageBuffer::from_vec(width, height, buf)
+            .expect("readback buffer was not big enough for image");
+        Some(imageops::flip_vertical(&img))
+    }
 }

src/window/wgpu_canvas.rs

@@ -104,6 +104,16 @@ enum PendingEvent {
     },
 }
 
+/// A GPU→CPU pixel readback still in flight (see `WgpuCanvas::begin_read_pixels`).
+struct PendingSnap {
+    buffer: wgpu::Buffer,
+    rx: std::sync::mpsc::Receiver<Result<(), wgpu::BufferAsyncError>>,
+    submission: wgpu::SubmissionIndex,
+    width: usize,
+    height: usize,
+    padded_bytes_per_row: usize,
+}
+
 /// A unified canvas based on wgpu that works on both native and web platforms.
 #[allow(dead_code)]
 pub struct WgpuCanvas {
@@ -126,6 +136,11 @@ pub struct WgpuCanvas {
     sample_count: u32,
     /// Texture for reading back pixels (for screenshots)
     readback_texture: wgpu::Texture,
+    /// Staging buffer reused across `read_pixels` calls, grown on demand, so
+    /// per-frame capture doesn't allocate (and free) a GPU buffer every call.
+    screenshot_staging: RefCell<Option<wgpu::Buffer>>,
+    /// Readback started by `begin_read_pixels`, completed by `finish_read_pixels`.
+    snap_pending: RefCell<Option<PendingSnap>>,
     /// Pending events from web callbacks (WASM only)
     #[cfg(target_arch = "wasm32")]
     pending_events: Rc<RefCell<Vec<WindowEvent>>>,
@@ -665,6 +680,8 @@ impl WgpuCanvas {
             msaa_view,
             sample_count,
             readback_texture,
+            screenshot_staging: RefCell::new(None),
+            snap_pending: RefCell::new(None),
             #[cfg(target_arch = "wasm32")]
             pending_events,
             #[cfg(target_arch = "wasm32")]
@@ -775,6 +792,8 @@ impl WgpuCanvas {
             msaa_view,
             sample_count,
             readback_texture,
+            screenshot_staging: RefCell::new(None),
+            snap_pending: RefCell::new(None),
             #[cfg(target_arch = "wasm32")]
             pending_events: Rc::new(RefCell::new(Vec::new())),
             #[cfg(target_arch = "wasm32")]
@@ -1310,6 +1329,25 @@ impl WgpuCanvas {
     /// Reads pixels from the readback texture into the provided buffer.
     /// Returns RGB data (3 bytes per pixel).
     pub fn read_pixels(&self, out: &mut Vec<u8>, x: usize, y: usize, width: usize, height: usize) {
+        self.begin_read_pixels(x, y, width, height);
+        self.finish_read_pixels(out);
+    }
+
+    /// Starts an asynchronous GPU→CPU readback of the readback texture and
+    /// returns immediately: it enqueues the texture→buffer copy and the buffer
+    /// map, but never waits on the GPU. Complete it — typically one frame
+    /// later, once the GPU has long finished the copy — with
+    /// [`Self::finish_read_pixels`]. Pipelining capture this way hides the
+    /// full CPU↔GPU sync that a blocking [`Self::read_pixels`] pays every
+    /// frame.
+    ///
+    /// A second `begin` before the previous readback was finished completes
+    /// and discards the previous one first (one readback in flight at a time).
+    pub fn begin_read_pixels(&self, x: usize, y: usize, width: usize, height: usize) {
+        if self.snap_pending.borrow().is_some() {
+            self.finish_read_pixels(&mut Vec::new());
+        }
+
         let ctxt = Context::get();
 
         // Calculate buffer size with alignment
@@ -1320,13 +1358,19 @@ impl WgpuCanvas {
         let padded_bytes_per_row = unpadded_bytes_per_row.div_ceil(align) * align;
         let buffer_size = padded_bytes_per_row * height;
 
-        // Create staging buffer
-        let staging_buffer = ctxt.create_buffer(&wgpu::BufferDescriptor {
-            label: Some("screenshot_staging_buffer"),
-            size: buffer_size as u64,
-            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
-            mapped_at_creation: false,
-        });
+        // Reuse the cached staging buffer when it is large enough; (re)create
+        // it otherwise. Capturing every frame (video export) then allocates
+        // exactly once instead of once per call.
+        let mut staging_slot = self.screenshot_staging.borrow_mut();
+        let staging_buffer = match staging_slot.take() {
+            Some(buffer) if buffer.size() >= buffer_size as u64 => buffer,
+            _ => ctxt.create_buffer(&wgpu::BufferDescriptor {
+                label: Some("screenshot_staging_buffer"),
+                size: buffer_size as u64,
+                usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
+                mapped_at_creation: false,
+            }),
+        };
 
         // Copy from readback texture to staging buffer
         let mut encoder = ctxt.create_command_encoder(Some("screenshot_copy_encoder"));
@@ -1357,20 +1401,55 @@ impl WgpuCanvas {
             },
         );
 
-        ctxt.submit(std::iter::once(encoder.finish()));
+        let submission = ctxt.submit_indexed(std::iter::once(encoder.finish()));
 
-        // Map the buffer and read the data
+        // Queue the map; completion is observed in `finish_read_pixels`.
         let buffer_slice = staging_buffer.slice(..);
         let (tx, rx) = std::sync::mpsc::channel();
         buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
             tx.send(result).unwrap();
         });
 
-        // Wait for the GPU to finish
-        let _ = ctxt.device.poll(wgpu::PollType::wait_indefinitely());
+        *self.snap_pending.borrow_mut() = Some(PendingSnap {
+            buffer: staging_buffer,
+            rx,
+            submission,
+            width,
+            height,
+            padded_bytes_per_row,
+        });
+    }
+
+    /// Completes a readback started by [`Self::begin_read_pixels`], filling
+    /// `out` with RGB data (3 bytes per pixel, rows bottom-to-top like
+    /// [`Self::read_pixels`]) and returning the captured `(width, height)`.
+    /// Returns `None` (and leaves `out` untouched) when no readback is in
+    /// flight. Blocks only until the copy's submission completes — a no-op
+    /// when a frame of GPU work has been submitted since the `begin`.
+    pub fn finish_read_pixels(&self, out: &mut Vec<u8>) -> Option<(u32, u32)> {
+        let PendingSnap {
+            buffer: staging_buffer,
+            rx,
+            submission,
+            width,
+            height,
+            padded_bytes_per_row,
+        } = self.snap_pending.borrow_mut().take()?;
+        let ctxt = Context::get();
+
+        // Wait only for the copy submission (and, transitively, the work it
+        // depends on) instead of polling the device indefinitely.
+        let _ = ctxt.device.poll(wgpu::PollType::Wait {
+            submission_index: Some(submission),
+            timeout: None,
+        });
         rx.recv().unwrap().unwrap();
 
+        let bytes_per_pixel = 4; // RGBA or BGRA
+        let unpadded_bytes_per_row = width * bytes_per_pixel;
+
         // Read the data
+        let buffer_slice = staging_buffer.slice(..);
         let data = buffer_slice.get_mapped_range();
 
         // Convert from BGRA/RGBA to RGB and handle row padding
@@ -1384,27 +1463,31 @@ impl WgpuCanvas {
         );
 
         // wgpu has origin at top-left, but we want bottom-left origin for OpenGL compatibility
-        // So we read rows in reverse order
+        // So we read rows in reverse order.
+        //
+        // The mapped range is uncached (write-combined) memory: scalar reads
+        // from it run at ~10 MB/s. memcpy each row into a cached local buffer
+        // first, then convert — orders of magnitude faster than indexing the
+        // mapped slice per byte.
+        let mut row_buf = vec![0u8; unpadded_bytes_per_row];
         for row in (0..height).rev() {
             let row_start = row * padded_bytes_per_row;
-            for col in 0..width {
-                let pixel_start = row_start + col * bytes_per_pixel;
+            row_buf.copy_from_slice(&data[row_start..row_start + unpadded_bytes_per_row]);
+            for px in row_buf.chunks_exact(bytes_per_pixel) {
                 if is_bgra {
-                    // BGRA -> RGB
-                    out.push(data[pixel_start + 2]); // R
-                    out.push(data[pixel_start + 1]); // G
-                    out.push(data[pixel_start]); // B
+                    out.extend_from_slice(&[px[2], px[1], px[0]]);
                 } else {
-                    // RGBA -> RGB
-                    out.push(data[pixel_start]); // R
-                    out.push(data[pixel_start + 1]); // G
-                    out.push(data[pixel_start + 2]); // B
+                    out.extend_from_slice(&[px[0], px[1], px[2]]);
                 }
             }
         }
 
         drop(data);
         staging_buffer.unmap();
+        // Return the buffer to the cache so the next `begin_read_pixels`
+        // (or blocking `read_pixels`) reuses it instead of allocating.
+        *self.screenshot_staging.borrow_mut() = Some(staging_buffer);
+        Some((width as u32, height as u32))
     }
 
     /// Gets the depth texture view for rendering.

src/window/window.rs

@@ -1050,6 +1050,16 @@ impl Window {
         usr_window
     }
 
+    /// Creates a headless window with custom setup options: a full-featured
+    /// [`Window`] backed by no OS window and no swapchain, rendering straight
+    /// into an off-screen texture. Unlike [`OffscreenSurface`](crate::window::OffscreenSurface)
+    /// this exposes the whole `Window` API (custom renderers, ray tracer,
+    /// `snap*` readbacks, …), for callers that drive the render loop
+    /// themselves and never present to a display.
+    pub async fn new_headless_with_setup(width: u32, height: u32, setup: CanvasSetup) -> Window {
+        Self::do_new_headless(width, height, Some(setup)).await
+    }
+
     /// Creates a headless window: a render target backed by no actual window,
     /// for off-screen rendering. Powers [`OffscreenSurface`](crate::window::OffscreenSurface).
     pub(super) async fn do_new_headless(