Add screenshot api (#7163)

Fixes https://github.com/bevyengine/bevy/issues/1207

# Objective

Right now, it's impossible to capture a screenshot of the entire window
without forking bevy. This is because
- The swapchain texture never has the COPY_SRC usage
- It can't be accessed without taking ownership of it
- Taking ownership of it breaks *a lot* of stuff

## Solution

- Introduce a dedicated api for taking a screenshot of a given bevy
window, and guarantee this screenshot will always match up with what
gets put on the screen.

---

## Changelog

- Added the `ScreenshotManager` resource with two functions,
`take_screenshot` and `save_screenshot_to_disk`

Cargo.toml

@@ -1909,6 +1909,16 @@ description = "Illustrates how to customize the default window settings"
 category = "Window"
 wasm = true
 
+[[example]]
+name = "screenshot"
+path = "examples/window/screenshot.rs"
+
+[package.metadata.example.screenshot]
+name = "Screenshot"
+description = "Shows how to save screenshots to disk"
+category = "Window"
+wasm = false
+
 [[example]]
 name = "transparent_window"
 path = "examples/window/transparent_window.rs"

crates/bevy_render/src/camera/camera.rs

@@ -423,7 +423,7 @@ impl NormalizedRenderTarget {
         match self {
             NormalizedRenderTarget::Window(window_ref) => windows
                 .get(&window_ref.entity())
-                .and_then(|window| window.swap_chain_texture.as_ref()),
+                .and_then(|window| window.swap_chain_texture_view.as_ref()),
             NormalizedRenderTarget::Image(image_handle) => {
                 images.get(image_handle).map(|image| &image.texture_view)
             }

crates/bevy_render/src/camera/camera_driver_node.rs

@@ -52,7 +52,7 @@ impl Node for CameraDriverNode {
                 continue;
             }
 
-            let Some(swap_chain_texture) = &window.swap_chain_texture else {
+            let Some(swap_chain_texture) = &window.swap_chain_texture_view else {
                 continue;
             };
 

crates/bevy_render/src/render_resource/texture.rs

@@ -51,31 +51,21 @@ define_atomic_id!(TextureViewId);
 render_resource_wrapper!(ErasedTextureView, wgpu::TextureView);
 render_resource_wrapper!(ErasedSurfaceTexture, wgpu::SurfaceTexture);
 
-/// This type combines wgpu's [`TextureView`](wgpu::TextureView) and
-/// [`SurfaceTexture`](wgpu::SurfaceTexture) into the same interface.
-#[derive(Clone, Debug)]
-pub enum TextureViewValue {
-    /// The value is an actual wgpu [`TextureView`](wgpu::TextureView).
-    TextureView(ErasedTextureView),
-
-    /// The value is a wgpu [`SurfaceTexture`](wgpu::SurfaceTexture), but dereferences to
-    /// a [`TextureView`](wgpu::TextureView).
-    SurfaceTexture {
-        // NOTE: The order of these fields is important because the view must be dropped before the
-        // frame is dropped
-        view: ErasedTextureView,
-        texture: ErasedSurfaceTexture,
-    },
-}
-
 /// Describes a [`Texture`] with its associated metadata required by a pipeline or [`BindGroup`](super::BindGroup).
-///
-/// May be converted from a [`TextureView`](wgpu::TextureView) or [`SurfaceTexture`](wgpu::SurfaceTexture)
-/// or dereferences to a wgpu [`TextureView`](wgpu::TextureView).
 #[derive(Clone, Debug)]
 pub struct TextureView {
     id: TextureViewId,
-    value: TextureViewValue,
+    value: ErasedTextureView,
+}
+
+pub struct SurfaceTexture {
+    value: ErasedSurfaceTexture,
+}
+
+impl SurfaceTexture {
+    pub fn try_unwrap(self) -> Option<wgpu::SurfaceTexture> {
+        self.value.try_unwrap()
+    }
 }
 
 impl TextureView {
@@ -84,34 +74,21 @@ impl TextureView {
     pub fn id(&self) -> TextureViewId {
         self.id
     }
-
-    /// Returns the [`SurfaceTexture`](wgpu::SurfaceTexture) of the texture view if it is of that type.
-    #[inline]
-    pub fn take_surface_texture(self) -> Option<wgpu::SurfaceTexture> {
-        match self.value {
-            TextureViewValue::TextureView(_) => None,
-            TextureViewValue::SurfaceTexture { texture, .. } => texture.try_unwrap(),
-        }
-    }
 }
 
 impl From<wgpu::TextureView> for TextureView {
     fn from(value: wgpu::TextureView) -> Self {
         TextureView {
             id: TextureViewId::new(),
-            value: TextureViewValue::TextureView(ErasedTextureView::new(value)),
+            value: ErasedTextureView::new(value),
         }
     }
 }
 
-impl From<wgpu::SurfaceTexture> for TextureView {
+impl From<wgpu::SurfaceTexture> for SurfaceTexture {
     fn from(value: wgpu::SurfaceTexture) -> Self {
-        let view = ErasedTextureView::new(value.texture.create_view(&Default::default()));
-        let texture = ErasedSurfaceTexture::new(value);
-
-        TextureView {
-            id: TextureViewId::new(),
-            value: TextureViewValue::SurfaceTexture { texture, view },
+        SurfaceTexture {
+            value: ErasedSurfaceTexture::new(value),
         }
     }
 }
@@ -121,10 +98,16 @@ impl Deref for TextureView {
 
     #[inline]
     fn deref(&self) -> &Self::Target {
-        match &self.value {
-            TextureViewValue::TextureView(value) => value,
-            TextureViewValue::SurfaceTexture { view, .. } => view,
-        }
+        &self.value
+    }
+}
+
+impl Deref for SurfaceTexture {
+    type Target = wgpu::SurfaceTexture;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.value
     }
 }
 

crates/bevy_render/src/renderer/graph_runner.rs

@@ -57,9 +57,12 @@ impl RenderGraphRunner {
         render_device: RenderDevice,
         queue: &wgpu::Queue,
         world: &World,
+        finalizer: impl FnOnce(&mut wgpu::CommandEncoder),
     ) -> Result<(), RenderGraphRunnerError> {
         let mut render_context = RenderContext::new(render_device);
         Self::run_graph(graph, None, &mut render_context, world, &[], None)?;
+        finalizer(render_context.command_encoder());
+
         {
             #[cfg(feature = "trace")]
             let _span = info_span!("submit_graph_commands").entered();

crates/bevy_render/src/renderer/mod.rs

@@ -35,6 +35,9 @@ pub fn render_system(world: &mut World) {
         render_device.clone(), // TODO: is this clone really necessary?
         &render_queue.0,
         world,
+        |encoder| {
+            crate::view::screenshot::submit_screenshot_commands(world, encoder);
+        },
     ) {
         error!("Error running render graph:");
         {
@@ -66,8 +69,8 @@ pub fn render_system(world: &mut World) {
 
         let mut windows = world.resource_mut::<ExtractedWindows>();
         for window in windows.values_mut() {
-            if let Some(texture_view) = window.swap_chain_texture.take() {
-                if let Some(surface_texture) = texture_view.take_surface_texture() {
+            if let Some(wrapped_texture) = window.swap_chain_texture.take() {
+                if let Some(surface_texture) = wrapped_texture.try_unwrap() {
                     surface_texture.present();
                 }
             }
@@ -81,6 +84,8 @@ pub fn render_system(world: &mut World) {
         );
     }
 
+    crate::view::screenshot::collect_screenshots(world);
+
     // update the time and send it to the app world
     let time_sender = world.resource::<TimeSender>();
     time_sender.0.try_send(Instant::now()).expect(

crates/bevy_render/src/texture/image_texture_conversion.rs

@@ -174,6 +174,7 @@ impl Image {
     /// - `TextureFormat::R8Unorm`
     /// - `TextureFormat::Rg8Unorm`
     /// - `TextureFormat::Rgba8UnormSrgb`
+    /// - `TextureFormat::Bgra8UnormSrgb`
     ///
     /// To convert [`Image`] to a different format see: [`Image::convert`].
     pub fn try_into_dynamic(self) -> anyhow::Result<DynamicImage> {
@@ -196,6 +197,20 @@ impl Image {
                 self.data,
             )
             .map(DynamicImage::ImageRgba8),
+            // This format is commonly used as the format for the swapchain texture
+            // This conversion is added here to support screenshots
+            TextureFormat::Bgra8UnormSrgb => ImageBuffer::from_raw(
+                self.texture_descriptor.size.width,
+                self.texture_descriptor.size.height,
+                {
+                    let mut data = self.data;
+                    for bgra in data.chunks_exact_mut(4) {
+                        bgra.swap(0, 2);
+                    }
+                    data
+                },
+            )
+            .map(DynamicImage::ImageRgba8),
             // Throw and error if conversion isn't supported
             texture_format => {
                 return Err(anyhow!(

crates/bevy_render/src/view/window.rs

@@ -1,6 +1,7 @@
 use crate::{
-    render_resource::TextureView,
+    render_resource::{PipelineCache, SpecializedRenderPipelines, SurfaceTexture, TextureView},
     renderer::{RenderAdapter, RenderDevice, RenderInstance},
+    texture::TextureFormatPixelInfo,
     Extract, ExtractSchedule, Render, RenderApp, RenderSet,
 };
 use bevy_app::{App, Plugin};
@@ -10,7 +11,13 @@ use bevy_window::{
     CompositeAlphaMode, PresentMode, PrimaryWindow, RawHandleWrapper, Window, WindowClosed,
 };
 use std::ops::{Deref, DerefMut};
-use wgpu::TextureFormat;
+use wgpu::{BufferUsages, TextureFormat, TextureUsages};
+
+pub mod screenshot;
+
+use screenshot::{
+    ScreenshotManager, ScreenshotPlugin, ScreenshotPreparedState, ScreenshotToScreenPipeline,
+};
 
 use super::Msaa;
 
@@ -27,10 +34,13 @@ pub enum WindowSystem {
 
 impl Plugin for WindowRenderPlugin {
     fn build(&self, app: &mut App) {
+        app.add_plugin(ScreenshotPlugin);
+
         if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
             render_app
                 .init_resource::<ExtractedWindows>()
                 .init_resource::<WindowSurfaces>()
+                .init_resource::<ScreenshotToScreenPipeline>()
                 .init_non_send_resource::<NonSendMarker>()
                 .add_systems(ExtractSchedule, extract_windows)
                 .configure_set(Render, WindowSystem::Prepare.in_set(RenderSet::Prepare))
@@ -46,11 +56,26 @@ pub struct ExtractedWindow {
     pub physical_width: u32,
     pub physical_height: u32,
     pub present_mode: PresentMode,
-    pub swap_chain_texture: Option<TextureView>,
+    /// Note: this will not always be the swap chain texture view. When taking a screenshot,
+    /// this will point to an alternative texture instead to allow for copying the render result
+    /// to CPU memory.
+    pub swap_chain_texture_view: Option<TextureView>,
+    pub swap_chain_texture: Option<SurfaceTexture>,
     pub swap_chain_texture_format: Option<TextureFormat>,
+    pub screenshot_memory: Option<ScreenshotPreparedState>,
     pub size_changed: bool,
     pub present_mode_changed: bool,
     pub alpha_mode: CompositeAlphaMode,
+    pub screenshot_func: Option<screenshot::ScreenshotFn>,
+}
+
+impl ExtractedWindow {
+    fn set_swapchain_texture(&mut self, frame: wgpu::SurfaceTexture) {
+        self.swap_chain_texture_view = Some(TextureView::from(
+            frame.texture.create_view(&Default::default()),
+        ));
+        self.swap_chain_texture = Some(SurfaceTexture::from(frame));
+    }
 }
 
 #[derive(Default, Resource)]
@@ -75,6 +100,7 @@ impl DerefMut for ExtractedWindows {
 
 fn extract_windows(
     mut extracted_windows: ResMut<ExtractedWindows>,
+    screenshot_manager: Extract<Res<ScreenshotManager>>,
     mut closed: Extract<EventReader<WindowClosed>>,
     windows: Extract<Query<(Entity, &Window, &RawHandleWrapper, Option<&PrimaryWindow>)>>,
 ) {
@@ -95,14 +121,17 @@ fn extract_windows(
             physical_height: new_height,
             present_mode: window.present_mode,
             swap_chain_texture: None,
+            swap_chain_texture_view: None,
             size_changed: false,
             swap_chain_texture_format: None,
             present_mode_changed: false,
             alpha_mode: window.composite_alpha_mode,
+            screenshot_func: None,
+            screenshot_memory: None,
         });
 
         // NOTE: Drop the swap chain frame here
-        extracted_window.swap_chain_texture = None;
+        extracted_window.swap_chain_texture_view = None;
         extracted_window.size_changed = new_width != extracted_window.physical_width
             || new_height != extracted_window.physical_height;
         extracted_window.present_mode_changed =
@@ -132,6 +161,15 @@ fn extract_windows(
     for closed_window in closed.iter() {
         extracted_windows.remove(&closed_window.window);
     }
+    // This lock will never block because `callbacks` is `pub(crate)` and this is the singular callsite where it's locked.
+    // Even if a user had multiple copies of this system, since the system has a mutable resource access the two systems would never run
+    // at the same time
+    // TODO: since this is guaranteed, should the lock be replaced with an UnsafeCell to remove the overhead, or is it minor enough to be ignored?
+    for (window, screenshot_func) in screenshot_manager.callbacks.lock().drain() {
+        if let Some(window) = extracted_windows.get_mut(&window) {
+            window.screenshot_func = Some(screenshot_func);
+        }
+    }
 }
 
 struct SurfaceData {
@@ -167,6 +205,7 @@ pub struct WindowSurfaces {
 ///   another alternative is to try to use [`ANGLE`](https://github.com/gfx-rs/wgpu#angle) and
 ///   [`Backends::GL`](crate::settings::Backends::GL) if your GPU/drivers support `OpenGL 4.3` / `OpenGL ES 3.0` or
 ///   later.
+#[allow(clippy::too_many_arguments)]
 pub fn prepare_windows(
     // By accessing a NonSend resource, we tell the scheduler to put this system on the main thread,
     // which is necessary for some OS s
@@ -176,6 +215,9 @@ pub fn prepare_windows(
     render_device: Res<RenderDevice>,
     render_instance: Res<RenderInstance>,
     render_adapter: Res<RenderAdapter>,
+    screenshot_pipeline: Res<ScreenshotToScreenPipeline>,
+    pipeline_cache: Res<PipelineCache>,
+    mut pipelines: ResMut<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>,
     mut msaa: ResMut<Msaa>,
 ) {
     for window in windows.windows.values_mut() {
@@ -285,18 +327,18 @@ pub fn prepare_windows(
             let frame = surface
                 .get_current_texture()
                 .expect("Error configuring surface");
-            window.swap_chain_texture = Some(TextureView::from(frame));
+            window.set_swapchain_texture(frame);
         } else {
             match surface.get_current_texture() {
                 Ok(frame) => {
-                    window.swap_chain_texture = Some(TextureView::from(frame));
+                    window.set_swapchain_texture(frame);
                 }
                 Err(wgpu::SurfaceError::Outdated) => {
                     render_device.configure_surface(surface, &surface_configuration);
                     let frame = surface
                         .get_current_texture()
                         .expect("Error reconfiguring surface");
-                    window.swap_chain_texture = Some(TextureView::from(frame));
+                    window.set_swapchain_texture(frame);
                 }
                 #[cfg(target_os = "linux")]
                 Err(wgpu::SurfaceError::Timeout) if may_erroneously_timeout() => {
@@ -311,5 +353,55 @@ pub fn prepare_windows(
             }
         };
         window.swap_chain_texture_format = Some(surface_data.format);
+
+        if window.screenshot_func.is_some() {
+            let texture = render_device.create_texture(&wgpu::TextureDescriptor {
+                label: Some("screenshot-capture-rendertarget"),
+                size: wgpu::Extent3d {
+                    width: surface_configuration.width,
+                    height: surface_configuration.height,
+                    depth_or_array_layers: 1,
+                },
+                mip_level_count: 1,
+                sample_count: 1,
+                dimension: wgpu::TextureDimension::D2,
+                format: surface_configuration.format,
+                usage: TextureUsages::RENDER_ATTACHMENT
+                    | TextureUsages::COPY_SRC
+                    | TextureUsages::TEXTURE_BINDING,
+                view_formats: &[],
+            });
+            let texture_view = texture.create_view(&Default::default());
+            let buffer = render_device.create_buffer(&wgpu::BufferDescriptor {
+                label: Some("screenshot-transfer-buffer"),
+                size: screenshot::get_aligned_size(
+                    window.physical_width,
+                    window.physical_height,
+                    surface_data.format.pixel_size() as u32,
+                ) as u64,
+                usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
+                mapped_at_creation: false,
+            });
+            let bind_group = render_device.create_bind_group(&wgpu::BindGroupDescriptor {
+                label: Some("screenshot-to-screen-bind-group"),
+                layout: &screenshot_pipeline.bind_group_layout,
+                entries: &[wgpu::BindGroupEntry {
+                    binding: 0,
+                    resource: wgpu::BindingResource::TextureView(&texture_view),
+                }],
+            });
+            let pipeline_id = pipelines.specialize(
+                &pipeline_cache,
+                &screenshot_pipeline,
+                surface_configuration.format,
+            );
+            window.swap_chain_texture_view = Some(texture_view);
+            window.screenshot_memory = Some(ScreenshotPreparedState {
+                texture,
+                buffer,
+                bind_group,
+                pipeline_id,
+            });
+        }
     }
 }

crates/bevy_render/src/view/window/screenshot.rs

@@ -0,0 +1,315 @@
+use std::{borrow::Cow, num::NonZeroU32, path::Path};
+
+use bevy_app::Plugin;
+use bevy_asset::{load_internal_asset, HandleUntyped};
+use bevy_ecs::prelude::*;
+use bevy_log::{error, info, info_span};
+use bevy_reflect::TypeUuid;
+use bevy_tasks::AsyncComputeTaskPool;
+use bevy_utils::HashMap;
+use parking_lot::Mutex;
+use thiserror::Error;
+use wgpu::{
+    CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT,
+};
+
+use crate::{
+    prelude::{Image, Shader},
+    render_resource::{
+        BindGroup, BindGroupLayout, Buffer, CachedRenderPipelineId, FragmentState, PipelineCache,
+        RenderPipelineDescriptor, SpecializedRenderPipeline, SpecializedRenderPipelines, Texture,
+        VertexState,
+    },
+    renderer::RenderDevice,
+    texture::TextureFormatPixelInfo,
+    RenderApp,
+};
+
+use super::ExtractedWindows;
+
+pub type ScreenshotFn = Box<dyn FnOnce(Image) + Send + Sync>;
+
+/// A resource which allows for taking screenshots of the window.
+#[derive(Resource, Default)]
+pub struct ScreenshotManager {
+    // this is in a mutex to enable extraction with only an immutable reference
+    pub(crate) callbacks: Mutex<HashMap<Entity, ScreenshotFn>>,
+}
+
+#[derive(Error, Debug)]
+#[error("A screenshot for this window has already been requested.")]
+pub struct ScreenshotAlreadyRequestedError;
+
+impl ScreenshotManager {
+    /// Signals the renderer to take a screenshot of this frame.
+    ///
+    /// The given callback will eventually be called on one of the [`AsyncComputeTaskPool`]s threads.
+    pub fn take_screenshot(
+        &mut self,
+        window: Entity,
+        callback: impl FnOnce(Image) + Send + Sync + 'static,
+    ) -> Result<(), ScreenshotAlreadyRequestedError> {
+        self.callbacks
+            .get_mut()
+            .try_insert(window, Box::new(callback))
+            .map(|_| ())
+            .map_err(|_| ScreenshotAlreadyRequestedError)
+    }
+
+    /// Signals the renderer to take a screenshot of this frame.
+    ///
+    /// The screenshot will eventually be saved to the given path, and the format will be derived from the extension.
+    pub fn save_screenshot_to_disk(
+        &mut self,
+        window: Entity,
+        path: impl AsRef<Path>,
+    ) -> Result<(), ScreenshotAlreadyRequestedError> {
+        let path = path.as_ref().to_owned();
+        self.take_screenshot(window, move |img| match img.try_into_dynamic() {
+            Ok(dyn_img) => match image::ImageFormat::from_path(&path) {
+                Ok(format) => {
+                    // discard the alpha channel which stores brightness values when HDR is enabled to make sure
+                    // the screenshot looks right
+                    let img = dyn_img.to_rgb8();
+                    match img.save_with_format(&path, format) {
+                        Ok(_) => info!("Screenshot saved to {}", path.display()),
+                        Err(e) => error!("Cannot save screenshot, IO error: {e}"),
+                    }
+                }
+                Err(e) => error!("Cannot save screenshot, requested format not recognized: {e}"),
+            },
+            Err(e) => error!("Cannot save screenshot, screen format cannot be understood: {e}"),
+        })
+    }
+}
+
+pub struct ScreenshotPlugin;
+
+const SCREENSHOT_SHADER_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 11918575842344596158);
+
+impl Plugin for ScreenshotPlugin {
+    fn build(&self, app: &mut bevy_app::App) {
+        app.init_resource::<ScreenshotManager>();
+
+        load_internal_asset!(
+            app,
+            SCREENSHOT_SHADER_HANDLE,
+            "screenshot.wgsl",
+            Shader::from_wgsl
+        );
+
+        if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
+            render_app.init_resource::<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>();
+        }
+    }
+}
+
+pub(crate) fn align_byte_size(value: u32) -> u32 {
+    value + (COPY_BYTES_PER_ROW_ALIGNMENT - (value % COPY_BYTES_PER_ROW_ALIGNMENT))
+}
+
+pub(crate) fn get_aligned_size(width: u32, height: u32, pixel_size: u32) -> u32 {
+    height * align_byte_size(width * pixel_size)
+}
+
+pub(crate) fn layout_data(width: u32, height: u32, format: TextureFormat) -> ImageDataLayout {
+    ImageDataLayout {
+        bytes_per_row: if height > 1 {
+            // 1 = 1 row
+            NonZeroU32::new(get_aligned_size(width, 1, format.pixel_size() as u32))
+        } else {
+            None
+        },
+        rows_per_image: None,
+        ..Default::default()
+    }
+}
+
+#[derive(Resource)]
+pub struct ScreenshotToScreenPipeline {
+    pub bind_group_layout: BindGroupLayout,
+}
+
+impl FromWorld for ScreenshotToScreenPipeline {
+    fn from_world(render_world: &mut World) -> Self {
+        let device = render_world.resource::<RenderDevice>();
+
+        let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
+            label: Some("screenshot-to-screen-bgl"),
+            entries: &[wgpu::BindGroupLayoutEntry {
+                binding: 0,
+                visibility: wgpu::ShaderStages::FRAGMENT,
+                ty: wgpu::BindingType::Texture {
+                    sample_type: wgpu::TextureSampleType::Float { filterable: false },
+                    view_dimension: wgpu::TextureViewDimension::D2,
+                    multisampled: false,
+                },
+                count: None,
+            }],
+        });
+
+        Self { bind_group_layout }
+    }
+}
+
+impl SpecializedRenderPipeline for ScreenshotToScreenPipeline {
+    type Key = TextureFormat;
+
+    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
+        RenderPipelineDescriptor {
+            label: Some(Cow::Borrowed("screenshot-to-screen")),
+            layout: vec![self.bind_group_layout.clone()],
+            vertex: VertexState {
+                buffers: vec![],
+                shader_defs: vec![],
+                entry_point: Cow::Borrowed("vs_main"),
+                shader: SCREENSHOT_SHADER_HANDLE.typed(),
+            },
+            primitive: wgpu::PrimitiveState {
+                topology: wgpu::PrimitiveTopology::TriangleList,
+                strip_index_format: None,
+                front_face: wgpu::FrontFace::Ccw,
+                cull_mode: Some(wgpu::Face::Back),
+                polygon_mode: wgpu::PolygonMode::Fill,
+                conservative: false,
+                unclipped_depth: false,
+            },
+            depth_stencil: None,
+            multisample: wgpu::MultisampleState::default(),
+            fragment: Some(FragmentState {
+                shader: SCREENSHOT_SHADER_HANDLE.typed(),
+                entry_point: Cow::Borrowed("fs_main"),
+                shader_defs: vec![],
+                targets: vec![Some(wgpu::ColorTargetState {
+                    format: key,
+                    blend: None,
+                    write_mask: wgpu::ColorWrites::ALL,
+                })],
+            }),
+            push_constant_ranges: Vec::new(),
+        }
+    }
+}
+
+pub struct ScreenshotPreparedState {
+    pub texture: Texture,
+    pub buffer: Buffer,
+    pub bind_group: BindGroup,
+    pub pipeline_id: CachedRenderPipelineId,
+}
+
+pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEncoder) {
+    let windows = world.resource::<ExtractedWindows>();
+    let pipelines = world.resource::<PipelineCache>();
+
+    for window in windows.values() {
+        if let Some(memory) = &window.screenshot_memory {
+            let width = window.physical_width;
+            let height = window.physical_height;
+            let texture_format = window.swap_chain_texture_format.unwrap();
+
+            encoder.copy_texture_to_buffer(
+                memory.texture.as_image_copy(),
+                wgpu::ImageCopyBuffer {
+                    buffer: &memory.buffer,
+                    layout: crate::view::screenshot::layout_data(width, height, texture_format),
+                },
+                Extent3d {
+                    width,
+                    height,
+                    ..Default::default()
+                },
+            );
+            if let Some(pipeline) = pipelines.get_render_pipeline(memory.pipeline_id) {
+                let true_swapchain_texture_view = window
+                    .swap_chain_texture
+                    .as_ref()
+                    .unwrap()
+                    .texture
+                    .create_view(&Default::default());
+                let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
+                    label: Some("screenshot_to_screen_pass"),
+                    color_attachments: &[Some(wgpu::RenderPassColorAttachment {
+                        view: &true_swapchain_texture_view,
+                        resolve_target: None,
+                        ops: wgpu::Operations {
+                            load: wgpu::LoadOp::Load,
+                            store: true,
+                        },
+                    })],
+                    depth_stencil_attachment: None,
+                });
+                pass.set_pipeline(pipeline);
+                pass.set_bind_group(0, &memory.bind_group, &[]);
+                pass.draw(0..3, 0..1);
+            }
+        }
+    }
+}
+
+pub(crate) fn collect_screenshots(world: &mut World) {
+    let _span = info_span!("collect_screenshots");
+
+    let mut windows = world.resource_mut::<ExtractedWindows>();
+    for window in windows.values_mut() {
+        if let Some(screenshot_func) = window.screenshot_func.take() {
+            let width = window.physical_width;
+            let height = window.physical_height;
+            let texture_format = window.swap_chain_texture_format.unwrap();
+            let pixel_size = texture_format.pixel_size();
+            let ScreenshotPreparedState { buffer, .. } = window.screenshot_memory.take().unwrap();
+
+            let finish = async move {
+                let (tx, rx) = async_channel::bounded(1);
+                let buffer_slice = buffer.slice(..);
+                // The polling for this map call is done every frame when the command queue is submitted.
+                buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
+                    let err = result.err();
+                    if err.is_some() {
+                        panic!("{}", err.unwrap().to_string());
+                    }
+                    tx.try_send(()).unwrap();
+                });
+                rx.recv().await.unwrap();
+                let data = buffer_slice.get_mapped_range();
+                // we immediately move the data to CPU memory to avoid holding the mapped view for long
+                let mut result = Vec::from(&*data);
+                drop(data);
+                drop(buffer);
+
+                if result.len() != ((width * height) as usize * pixel_size) {
+                    // Our buffer has been padded because we needed to align to a multiple of 256.
+                    // We remove this padding here
+                    let initial_row_bytes = width as usize * pixel_size;
+                    let buffered_row_bytes = align_byte_size(width * pixel_size as u32) as usize;
+
+                    let mut take_offset = buffered_row_bytes;
+                    let mut place_offset = initial_row_bytes;
+                    for _ in 1..height {
+                        result.copy_within(
+                            take_offset..take_offset + buffered_row_bytes,
+                            place_offset,
+                        );
+                        take_offset += buffered_row_bytes;
+                        place_offset += initial_row_bytes;
+                    }
+                    result.truncate(initial_row_bytes * height as usize);
+                }
+
+                screenshot_func(Image::new(
+                    Extent3d {
+                        width,
+                        height,
+                        depth_or_array_layers: 1,
+                    },
+                    wgpu::TextureDimension::D2,
+                    result,
+                    texture_format,
+                ));
+            };
+
+            AsyncComputeTaskPool::get().spawn(finish).detach();
+        }
+    }
+}

crates/bevy_render/src/view/window/screenshot.wgsl

@@ -0,0 +1,16 @@
+// This vertex shader will create a triangle that will cover the entire screen
+// with minimal effort, avoiding the need for a vertex buffer etc.
+@vertex
+fn vs_main(@builtin(vertex_index) in_vertex_index: u32) -> @builtin(position) vec4<f32> {
+    let x = f32((in_vertex_index & 1u) << 2u);
+    let y = f32((in_vertex_index & 2u) << 1u);
+    return vec4<f32>(x - 1.0, y - 1.0, 0.0, 1.0);
+}
+
+@group(0) @binding(0) var t: texture_2d<f32>;
+
+@fragment
+fn fs_main(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> {
+    let coords = floor(pos.xy);
+    return textureLoad(t, vec2<i32>(coords), 0i);
+}

examples/README.md

@@ -353,6 +353,7 @@ Example | Description
 [Low Power](../examples/window/low_power.rs) | Demonstrates settings to reduce power use for bevy applications
 [Multiple Windows](../examples/window/multiple_windows.rs) | Demonstrates creating multiple windows, and rendering to them
 [Scale Factor Override](../examples/window/scale_factor_override.rs) | Illustrates how to customize the default window settings
+[Screenshot](../examples/window/screenshot.rs) | Shows how to save screenshots to disk
 [Transparent Window](../examples/window/transparent_window.rs) | Illustrates making the window transparent and hiding the window decoration
 [Window Resizing](../examples/window/window_resizing.rs) | Demonstrates resizing and responding to resizing a window
 [Window Settings](../examples/window/window_settings.rs) | Demonstrates customizing default window settings

examples/window/screenshot.rs

@@ -0,0 +1,64 @@
+//! An example showing how to save screenshots to disk
+
+use bevy::prelude::*;
+use bevy::render::view::screenshot::ScreenshotManager;
+use bevy::window::PrimaryWindow;
+
+fn main() {
+    App::new()
+        .add_plugins(DefaultPlugins)
+        .add_systems(Startup, setup)
+        .add_systems(Update, screenshot_on_f12)
+        .run();
+}
+
+fn screenshot_on_f12(
+    input: Res<Input<KeyCode>>,
+    main_window: Query<Entity, With<PrimaryWindow>>,
+    mut screenshot_manager: ResMut<ScreenshotManager>,
+    mut counter: Local<u32>,
+) {
+    if input.just_pressed(KeyCode::F12) {
+        let path = format!("./screenshot-{}.png", *counter);
+        *counter += 1;
+        screenshot_manager
+            .save_screenshot_to_disk(main_window.single(), path)
+            .unwrap();
+    }
+}
+
+/// set up a simple 3D scene
+fn setup(
+    mut commands: Commands,
+    mut meshes: ResMut<Assets<Mesh>>,
+    mut materials: ResMut<Assets<StandardMaterial>>,
+) {
+    // plane
+    commands.spawn(PbrBundle {
+        mesh: meshes.add(shape::Plane::from_size(5.0).into()),
+        material: materials.add(Color::rgb(0.3, 0.5, 0.3).into()),
+        ..default()
+    });
+    // cube
+    commands.spawn(PbrBundle {
+        mesh: meshes.add(Mesh::from(shape::Cube { size: 1.0 })),
+        material: materials.add(Color::rgb(0.8, 0.7, 0.6).into()),
+        transform: Transform::from_xyz(0.0, 0.5, 0.0),
+        ..default()
+    });
+    // light
+    commands.spawn(PointLightBundle {
+        point_light: PointLight {
+            intensity: 1500.0,
+            shadows_enabled: true,
+            ..default()
+        },
+        transform: Transform::from_xyz(4.0, 8.0, 4.0),
+        ..default()
+    });
+    // camera
+    commands.spawn(Camera3dBundle {
+        transform: Transform::from_xyz(-2.0, 2.5, 5.0).looking_at(Vec3::ZERO, Vec3::Y),
+        ..default()
+    });
+}