Rewrite screenshots. (#14833)

# Objective

Rewrite screenshotting to be able to accept any `RenderTarget`.

Closes #12478 

## Solution

Previously, screenshotting relied on setting a variety of state on the
requested window. When extracted, the window's `swap_chain_texture_view`
property would be swapped out with a texture_view created that frame for
the screenshot pipeline to write back to the cpu.

Besides being tightly coupled to window in a way that prevented
screenshotting other render targets, this approach had the drawback of
relying on the implicit state of `swap_chain_texture_view` being
returned from a `NormalizedRenderTarget` when view targets were
prepared. Because property is set every frame for windows, that wasn't a
problem, but poses a problem for render target images. Namely, to do the
equivalent trick, we'd have to replace the `GpuImage`'s texture view,
and somehow restore it later.

As such, this PR creates a new `prepare_view_textures` system which runs
before `prepare_view_targets` that allows a new `prepare_screenshots`
system to be sandwiched between and overwrite the render targets texture
view if a screenshot has been requested that frame for the given target.

Additionally, screenshotting itself has been changed to use a component
+ observer pattern. We now spawn a `Screenshot` component into the
world, whose lifetime is tracked with a series of marker components.
When the screenshot is read back to the CPU, we send the image over a
channel back to the main world where an observer fires on the screenshot
entity before being despawned the next frame. This allows the user to
access resources in their save callback that might be useful (e.g.
uploading the screenshot over the network, etc.).

## Testing


![image](https://github.com/user-attachments/assets/48f19aed-d9e1-4058-bb17-82b37f992b7b)


TODO:
- [x] Web
- [ ] Manual texture view

---

## Showcase

render to texture example:
<img
src="https://github.com/user-attachments/assets/612ac47b-8a24-4287-a745-3051837963b0"
width=200/>

web saving still works:
<img
src="https://github.com/user-attachments/assets/e2a15b17-1ff5-4006-ab2a-e5cc74888b9c"
width=200/>

## Migration Guide

`ScreenshotManager` has been removed. To take a screenshot, spawn a
`Screenshot` entity with the specified render target and provide an
observer targeting the `ScreenshotCaptured` event. See the
`window/screenshot` example to see an example.

---------

Co-authored-by: Kristoffer Søholm <k.soeholm@gmail.com>

crates/bevy_dev_tools/src/ci_testing/mod.rs

@@ -7,6 +7,7 @@ pub use self::config::*;
 
 use bevy_app::prelude::*;
 use bevy_ecs::schedule::IntoSystemConfigs;
+use bevy_render::view::screenshot::trigger_screenshots;
 use bevy_time::TimeUpdateStrategy;
 use std::time::Duration;
 
@@ -51,7 +52,9 @@ impl Plugin for CiTestingPlugin {
             .insert_resource(config)
             .add_systems(
                 Update,
-                systems::send_events.before(bevy_window::close_when_requested),
+                systems::send_events
+                    .before(trigger_screenshots)
+                    .before(bevy_window::close_when_requested),
             );
     }
 }

crates/bevy_dev_tools/src/ci_testing/systems.rs

@@ -1,9 +1,8 @@
 use super::config::*;
 use bevy_app::AppExit;
 use bevy_ecs::prelude::*;
-use bevy_render::view::screenshot::ScreenshotManager;
+use bevy_render::view::screenshot::{save_to_disk, Screenshot};
-use bevy_utils::tracing::{debug, info, warn};
+use bevy_utils::tracing::{debug, info};
-use bevy_window::PrimaryWindow;
 
 pub(crate) fn send_events(world: &mut World, mut current_frame: Local<u32>) {
     let mut config = world.resource_mut::<CiTestingConfig>();
@@ -23,21 +22,10 @@ pub(crate) fn send_events(world: &mut World, mut current_frame: Local<u32>) {
                 info!("Exiting after {} frames. Test successful!", *current_frame);
             }
             CiTestingEvent::Screenshot => {
-                let mut primary_window_query =
-                    world.query_filtered::<Entity, With<PrimaryWindow>>();
-                let Ok(main_window) = primary_window_query.get_single(world) else {
-                    warn!("Requesting screenshot, but PrimaryWindow is not available");
-                    continue;
-                };
-                let Some(mut screenshot_manager) = world.get_resource_mut::<ScreenshotManager>()
-                else {
-                    warn!("Requesting screenshot, but ScreenshotManager is not available");
-                    continue;
-                };
                 let path = format!("./screenshot-{}.png", *current_frame);
-                screenshot_manager
+                world
-                    .save_screenshot_to_disk(main_window, path)
+                    .spawn(Screenshot::primary_window())
-                    .unwrap();
+                    .observe(save_to_disk(path));
                 info!("Took a screenshot at frame {}.", *current_frame);
             }
             // Custom events are forwarded to the world.

crates/bevy_render/src/view/mod.rs

@@ -5,6 +5,7 @@ use bevy_asset::{load_internal_asset, Handle};
 pub use visibility::*;
 pub use window::*;
 
+use crate::camera::NormalizedRenderTarget;
 use crate::extract_component::ExtractComponentPlugin;
 use crate::{
     camera::{
@@ -25,12 +26,13 @@ use crate::{
 };
 use bevy_app::{App, Plugin};
 use bevy_color::LinearRgba;
+use bevy_derive::{Deref, DerefMut};
 use bevy_ecs::prelude::*;
 use bevy_math::{mat3, vec2, vec3, Mat3, Mat4, UVec4, Vec2, Vec3, Vec4, Vec4Swizzles};
 use bevy_reflect::{std_traits::ReflectDefault, Reflect};
 use bevy_render_macros::ExtractComponent;
 use bevy_transform::components::GlobalTransform;
-use bevy_utils::HashMap;
+use bevy_utils::{hashbrown::hash_map::Entry, HashMap};
 use std::{
     ops::Range,
     sync::{
@@ -119,6 +121,9 @@ impl Plugin for ViewPlugin {
             render_app.add_systems(
                 Render,
                 (
+                    prepare_view_attachments
+                        .in_set(RenderSet::ManageViews)
+                        .before(prepare_view_targets),
                     prepare_view_targets
                         .in_set(RenderSet::ManageViews)
                         .after(prepare_windows)
@@ -132,7 +137,9 @@ impl Plugin for ViewPlugin {
 
     fn finish(&self, app: &mut App) {
         if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
-            render_app.init_resource::<ViewUniforms>();
+            render_app
+                .init_resource::<ViewUniforms>()
+                .init_resource::<ViewTargetAttachments>();
         }
     }
 }
@@ -458,6 +465,13 @@ pub struct ViewTarget {
     out_texture: OutputColorAttachment,
 }
 
+/// Contains [`OutputColorAttachment`] used for each target present on any view in the current
+/// frame, after being prepared by [`prepare_view_attachments`]. Users that want to override
+/// the default output color attachment for a specific target can do so by adding a
+/// [`OutputColorAttachment`] to this resource before [`prepare_view_targets`] is called.
+#[derive(Resource, Default, Deref, DerefMut)]
+pub struct ViewTargetAttachments(HashMap<NormalizedRenderTarget, OutputColorAttachment>);
+
 pub struct PostProcessWrite<'a> {
     pub source: &'a TextureView,
     pub destination: &'a TextureView,
@@ -794,11 +808,41 @@ struct MainTargetTextures {
     main_texture: Arc<AtomicUsize>,
 }
 
-#[allow(clippy::too_many_arguments)]
+/// Prepares the view target [`OutputColorAttachment`] for each view in the current frame.
-pub fn prepare_view_targets(
+pub fn prepare_view_attachments(
-    mut commands: Commands,
     windows: Res<ExtractedWindows>,
     images: Res<RenderAssets<GpuImage>>,
+    manual_texture_views: Res<ManualTextureViews>,
+    cameras: Query<&ExtractedCamera>,
+    mut view_target_attachments: ResMut<ViewTargetAttachments>,
+) {
+    view_target_attachments.clear();
+    for camera in cameras.iter() {
+        let Some(target) = &camera.target else {
+            continue;
+        };
+
+        match view_target_attachments.entry(target.clone()) {
+            Entry::Occupied(_) => {}
+            Entry::Vacant(entry) => {
+                let Some(attachment) = target
+                    .get_texture_view(&windows, &images, &manual_texture_views)
+                    .cloned()
+                    .zip(target.get_texture_format(&windows, &images, &manual_texture_views))
+                    .map(|(view, format)| {
+                        OutputColorAttachment::new(view.clone(), format.add_srgb_suffix())
+                    })
+                else {
+                    continue;
+                };
+                entry.insert(attachment);
+            }
+        };
+    }
+}
+
+pub fn prepare_view_targets(
+    mut commands: Commands,
     clear_color_global: Res<ClearColor>,
     render_device: Res<RenderDevice>,
     mut texture_cache: ResMut<TextureCache>,
@@ -809,24 +853,16 @@ pub fn prepare_view_targets(
         &CameraMainTextureUsages,
         &Msaa,
     )>,
-    manual_texture_views: Res<ManualTextureViews>,
+    view_target_attachments: Res<ViewTargetAttachments>,
 ) {
     let mut textures = HashMap::default();
-    let mut output_textures = HashMap::default();
     for (entity, camera, view, texture_usage, msaa) in cameras.iter() {
         let (Some(target_size), Some(target)) = (camera.physical_target_size, &camera.target)
         else {
             continue;
         };
 
-        let Some(out_texture) = output_textures.entry(target.clone()).or_insert_with(|| {
+        let Some(out_attachment) = view_target_attachments.get(target) else {
-            target
-                .get_texture_view(&windows, &images, &manual_texture_views)
-                .zip(target.get_texture_format(&windows, &images, &manual_texture_views))
-                .map(|(view, format)| {
-                    OutputColorAttachment::new(view.clone(), format.add_srgb_suffix())
-                })
-        }) else {
             continue;
         };
 
@@ -913,7 +949,7 @@ pub fn prepare_view_targets(
             main_texture: main_textures.main_texture.clone(),
             main_textures,
             main_texture_format,
-            out_texture: out_texture.clone(),
+            out_texture: out_attachment.clone(),
         });
     }
 }

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

@@ -1,9 +1,6 @@
 use crate::{
-    render_resource::{
+    render_resource::{SurfaceTexture, TextureView},
-        BindGroupEntries, PipelineCache, SpecializedRenderPipelines, SurfaceTexture, TextureView,
-    },
     renderer::{RenderAdapter, RenderDevice, RenderInstance},
-    texture::TextureFormatPixelInfo,
     Extract, ExtractSchedule, Render, RenderApp, RenderSet, WgpuWrapper,
 };
 use bevy_app::{App, Last, Plugin};
@@ -18,21 +15,16 @@ use bevy_winit::CustomCursorCache;
 use std::{
     num::NonZeroU32,
     ops::{Deref, DerefMut},
-    sync::PoisonError,
 };
 use wgpu::{
-    BufferUsages, SurfaceConfiguration, SurfaceTargetUnsafe, TextureFormat, TextureUsages,
+    SurfaceConfiguration, SurfaceTargetUnsafe, TextureFormat, TextureUsages, TextureViewDescriptor,
-    TextureViewDescriptor,
 };
 
 pub mod cursor;
 pub mod screenshot;
 
-use screenshot::{
-    ScreenshotManager, ScreenshotPlugin, ScreenshotPreparedState, ScreenshotToScreenPipeline,
-};
-
 use self::cursor::update_cursors;
+use screenshot::{ScreenshotPlugin, ScreenshotToScreenPipeline};
 
 pub struct WindowRenderPlugin;
 
@@ -78,11 +70,9 @@ pub struct ExtractedWindow {
     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 {
@@ -120,7 +110,6 @@ impl DerefMut for ExtractedWindows {
 
 fn extract_windows(
     mut extracted_windows: ResMut<ExtractedWindows>,
-    screenshot_manager: Extract<Res<ScreenshotManager>>,
     mut closing: Extract<EventReader<WindowClosing>>,
     windows: Extract<Query<(Entity, &Window, &RawHandleWrapper, Option<&PrimaryWindow>)>>,
     mut removed: Extract<RemovedComponents<RawHandleWrapper>>,
@@ -149,8 +138,6 @@ fn extract_windows(
             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
@@ -189,20 +176,6 @@ fn extract_windows(
         extracted_windows.remove(&removed_window);
         window_surfaces.remove(&removed_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()
-        .unwrap_or_else(PoisonError::into_inner)
-        .drain()
-    {
-        if let Some(window) = extracted_windows.get_mut(&window) {
-            window.screenshot_func = Some(screenshot_func);
-        }
-    }
 }
 
 struct SurfaceData {
@@ -254,9 +227,6 @@ pub fn prepare_windows(
     mut windows: ResMut<ExtractedWindows>,
     mut window_surfaces: ResMut<WindowSurfaces>,
     render_device: Res<RenderDevice>,
-    screenshot_pipeline: Res<ScreenshotToScreenPipeline>,
-    pipeline_cache: Res<PipelineCache>,
-    mut pipelines: ResMut<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>,
     #[cfg(target_os = "linux")] render_instance: Res<RenderInstance>,
 ) {
     for window in windows.windows.values_mut() {
@@ -340,53 +310,6 @@ pub fn prepare_windows(
             }
         };
         window.swap_chain_texture_format = Some(surface_data.configuration.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_data.configuration.width,
-                    height: surface_data.configuration.height,
-                    depth_or_array_layers: 1,
-                },
-                mip_level_count: 1,
-                sample_count: 1,
-                dimension: wgpu::TextureDimension::D2,
-                format: surface_data.configuration.format.add_srgb_suffix(),
-                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.configuration.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(
-                "screenshot-to-screen-bind-group",
-                &screenshot_pipeline.bind_group_layout,
-                &BindGroupEntries::single(&texture_view),
-            );
-            let pipeline_id = pipelines.specialize(
-                &pipeline_cache,
-                &screenshot_pipeline,
-                surface_data.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

@@ -1,16 +1,13 @@
-use std::{borrow::Cow, path::Path, sync::PoisonError};
+use super::ExtractedWindows;
-
+use crate::camera::{
-use bevy_app::Plugin;
+    ManualTextureViewHandle, ManualTextureViews, NormalizedRenderTarget, RenderTarget,
-use bevy_asset::{load_internal_asset, Handle};
+};
-use bevy_ecs::{entity::EntityHashMap, prelude::*};
+use crate::render_asset::RenderAssets;
-use bevy_tasks::AsyncComputeTaskPool;
+use crate::render_resource::{BindGroupEntries, BufferUsages, TextureUsages, TextureView};
-use bevy_utils::tracing::{error, info, info_span};
+use crate::texture::{GpuImage, OutputColorAttachment};
-use std::sync::Mutex;
+use crate::view::{
-use thiserror::Error;
+    prepare_view_attachments, prepare_view_targets, ViewTargetAttachments, WindowSurfaces,
-use wgpu::{
-    CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT,
 };
-
 use crate::{
     prelude::{Image, Shader},
     render_asset::RenderAssetUsages,
@@ -21,51 +18,116 @@ use crate::{
     },
     renderer::RenderDevice,
     texture::TextureFormatPixelInfo,
-    RenderApp,
+    ExtractSchedule, MainWorld, Render, RenderApp, RenderSet,
+};
+use bevy_app::{First, Plugin, Update};
+use bevy_asset::{load_internal_asset, Handle};
+use bevy_derive::{Deref, DerefMut};
+use bevy_ecs::event::event_update_system;
+use bevy_ecs::system::SystemState;
+use bevy_ecs::{entity::EntityHashMap, prelude::*};
+use bevy_hierarchy::DespawnRecursiveExt;
+use bevy_reflect::Reflect;
+use bevy_tasks::AsyncComputeTaskPool;
+use bevy_utils::tracing::{error, info, warn};
+use bevy_utils::{default, HashSet};
+use bevy_window::{PrimaryWindow, WindowRef};
+use std::ops::Deref;
+use std::sync::mpsc::{Receiver, Sender};
+use std::sync::{Arc, Mutex};
+use std::{borrow::Cow, path::Path};
+use wgpu::{
+    CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT,
 };
 
-use super::ExtractedWindows;
+#[derive(Event, Deref, DerefMut, Reflect, Debug)]
+#[reflect(Debug)]
+pub struct ScreenshotCaptured(pub Image);
 
-pub type ScreenshotFn = Box<dyn FnOnce(Image) + Send + Sync>;
+/// A component that signals to the renderer to capture a screenshot this frame.
+///
+/// This component should be spawned on a new entity with an observer that will trigger
+/// with [`ScreenshotCaptured`] when the screenshot is ready.
+///
+/// Screenshots are captured asynchronously and may not be available immediately after the frame
+/// that the component is spawned on. The observer should be used to handle the screenshot when it
+/// is ready.
+///
+/// Note that the screenshot entity will be despawned after the screenshot is captured and the
+/// observer is triggered.
+///
+/// # Usage
+///
+/// ```
+/// # use bevy_ecs::prelude::*;
+/// # use bevy_render::view::screenshot::{save_to_disk, Screenshot};
+///
+/// fn take_screenshot(mut commands: Commands) {
+///    commands.spawn(Screenshot::primary_window())
+///       .observe(save_to_disk("screenshot.png"));
+/// }
+/// ```
+#[derive(Component, Deref, DerefMut, Reflect, Debug)]
+#[reflect(Component, Debug)]
+pub struct Screenshot(pub RenderTarget);
 
-/// A resource which allows for taking screenshots of the window.
+/// A marker component that indicates that a screenshot is currently being captured.
-#[derive(Resource, Default)]
+#[derive(Component)]
-pub struct ScreenshotManager {
+pub struct Capturing;
-    // this is in a mutex to enable extraction with only an immutable reference
+
-    pub(crate) callbacks: Mutex<EntityHashMap<ScreenshotFn>>,
+/// A marker component that indicates that a screenshot has been captured, the image is ready, and
+/// the screenshot entity can be despawned.
+#[derive(Component)]
+pub struct Captured;
+
+impl Screenshot {
+    /// Capture a screenshot of the provided window entity.
+    pub fn window(window: Entity) -> Self {
+        Self(RenderTarget::Window(WindowRef::Entity(window)))
     }
 
-#[derive(Error, Debug)]
+    /// Capture a screenshot of the primary window, if one exists.
-#[error("A screenshot for this window has already been requested.")]
+    pub fn primary_window() -> Self {
-pub struct ScreenshotAlreadyRequestedError;
+        Self(RenderTarget::Window(WindowRef::Primary))
-
-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()
-            .unwrap_or_else(PoisonError::into_inner)
-            .try_insert(window, Box::new(callback))
-            .map(|_| ())
-            .map_err(|_| ScreenshotAlreadyRequestedError)
     }
 
-    /// Signals the renderer to take a screenshot of this frame.
+    /// Capture a screenshot of the provided render target image.
-    ///
+    pub fn image(image: Handle<Image>) -> Self {
-    /// The screenshot will eventually be saved to the given path, and the format will be derived from the extension.
+        Self(RenderTarget::Image(image))
-    pub fn save_screenshot_to_disk(
+    }
-        &mut self,
+
-        window: Entity,
+    /// Capture a screenshot of the provided manual texture view.
-        path: impl AsRef<Path>,
+    pub fn texture_view(texture_view: ManualTextureViewHandle) -> Self {
-    ) -> Result<(), ScreenshotAlreadyRequestedError> {
+        Self(RenderTarget::TextureView(texture_view))
+    }
+}
+
+struct ScreenshotPreparedState {
+    pub texture: Texture,
+    pub buffer: Buffer,
+    pub bind_group: BindGroup,
+    pub pipeline_id: CachedRenderPipelineId,
+    pub size: Extent3d,
+}
+
+#[derive(Resource, Deref, DerefMut)]
+pub struct CapturedScreenshots(pub Arc<Mutex<Receiver<(Entity, Image)>>>);
+
+#[derive(Resource, Deref, DerefMut, Default)]
+struct RenderScreenshotTargets(EntityHashMap<NormalizedRenderTarget>);
+
+#[derive(Resource, Deref, DerefMut, Default)]
+struct RenderScreenshotsPrepared(EntityHashMap<ScreenshotPreparedState>);
+
+#[derive(Resource, Deref, DerefMut)]
+struct RenderScreenshotsSender(Sender<(Entity, Image)>);
+
+/// Saves the captured screenshot to disk at the provided path.
+pub fn save_to_disk(path: impl AsRef<Path>) -> impl FnMut(Trigger<ScreenshotCaptured>) {
     let path = path.as_ref().to_owned();
-        self.take_screenshot(window, move |img| match img.try_into_dynamic() {
+    move |trigger| {
+        let img = trigger.event().deref().clone();
+        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
@@ -118,9 +180,222 @@ impl ScreenshotManager {
                 Err(e) => error!("Cannot save screenshot, requested format not recognized: {e}"),
             },
             Err(e) => error!("Cannot save screenshot, screen format cannot be understood: {e}"),
-        })
         }
     }
+}
+
+fn clear_screenshots(mut commands: Commands, screenshots: Query<Entity, With<Captured>>) {
+    for entity in screenshots.iter() {
+        commands.entity(entity).despawn_recursive();
+    }
+}
+
+pub fn trigger_screenshots(
+    mut commands: Commands,
+    captured_screenshots: ResMut<CapturedScreenshots>,
+) {
+    let captured_screenshots = captured_screenshots.lock().unwrap();
+    while let Ok((entity, image)) = captured_screenshots.try_recv() {
+        commands.entity(entity).insert(Captured);
+        commands.trigger_targets(ScreenshotCaptured(image), entity);
+    }
+}
+
+fn extract_screenshots(
+    mut targets: ResMut<RenderScreenshotTargets>,
+    mut main_world: ResMut<MainWorld>,
+    mut system_state: Local<
+        Option<
+            SystemState<(
+                Commands,
+                Query<Entity, With<PrimaryWindow>>,
+                Query<(Entity, &Screenshot), Without<Capturing>>,
+            )>,
+        >,
+    >,
+    mut seen_targets: Local<HashSet<NormalizedRenderTarget>>,
+) {
+    if system_state.is_none() {
+        *system_state = Some(SystemState::new(&mut main_world));
+    }
+    let system_state = system_state.as_mut().unwrap();
+    let (mut commands, primary_window, screenshots) = system_state.get_mut(&mut main_world);
+
+    targets.clear();
+    seen_targets.clear();
+
+    let primary_window = primary_window.iter().next();
+
+    for (entity, screenshot) in screenshots.iter() {
+        let render_target = screenshot.0.clone();
+        let Some(render_target) = render_target.normalize(primary_window) else {
+            warn!(
+                "Unknown render target for screenshot, skipping: {:?}",
+                render_target
+            );
+            continue;
+        };
+        if seen_targets.contains(&render_target) {
+            warn!(
+                "Duplicate render target for screenshot, skipping entity {:?}: {:?}",
+                entity, render_target
+            );
+            // If we don't despawn the entity here, it will be captured again in the next frame
+            commands.entity(entity).despawn_recursive();
+            continue;
+        }
+        seen_targets.insert(render_target.clone());
+        targets.insert(entity, render_target);
+        commands.entity(entity).insert(Capturing);
+    }
+
+    system_state.apply(&mut main_world);
+}
+
+#[allow(clippy::too_many_arguments)]
+fn prepare_screenshots(
+    targets: Res<RenderScreenshotTargets>,
+    mut prepared: ResMut<RenderScreenshotsPrepared>,
+    window_surfaces: Res<WindowSurfaces>,
+    render_device: Res<RenderDevice>,
+    screenshot_pipeline: Res<ScreenshotToScreenPipeline>,
+    pipeline_cache: Res<PipelineCache>,
+    mut pipelines: ResMut<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>,
+    images: Res<RenderAssets<GpuImage>>,
+    manual_texture_views: Res<ManualTextureViews>,
+    mut view_target_attachments: ResMut<ViewTargetAttachments>,
+) {
+    prepared.clear();
+    for (entity, target) in targets.iter() {
+        match target {
+            NormalizedRenderTarget::Window(window) => {
+                let window = window.entity();
+                let Some(surface_data) = window_surfaces.surfaces.get(&window) else {
+                    warn!("Unknown window for screenshot, skipping: {:?}", window);
+                    continue;
+                };
+                let format = surface_data.configuration.format.add_srgb_suffix();
+                let size = Extent3d {
+                    width: surface_data.configuration.width,
+                    height: surface_data.configuration.height,
+                    ..default()
+                };
+                let (texture_view, state) = prepare_screenshot_state(
+                    size,
+                    format,
+                    &render_device,
+                    &screenshot_pipeline,
+                    &pipeline_cache,
+                    &mut pipelines,
+                );
+                prepared.insert(*entity, state);
+                view_target_attachments.insert(
+                    target.clone(),
+                    OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()),
+                );
+            }
+            NormalizedRenderTarget::Image(image) => {
+                let Some(gpu_image) = images.get(image) else {
+                    warn!("Unknown image for screenshot, skipping: {:?}", image);
+                    continue;
+                };
+                let format = gpu_image.texture_format;
+                let size = Extent3d {
+                    width: gpu_image.size.x,
+                    height: gpu_image.size.y,
+                    ..default()
+                };
+                let (texture_view, state) = prepare_screenshot_state(
+                    size,
+                    format,
+                    &render_device,
+                    &screenshot_pipeline,
+                    &pipeline_cache,
+                    &mut pipelines,
+                );
+                prepared.insert(*entity, state);
+                view_target_attachments.insert(
+                    target.clone(),
+                    OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()),
+                );
+            }
+            NormalizedRenderTarget::TextureView(texture_view) => {
+                let Some(manual_texture_view) = manual_texture_views.get(texture_view) else {
+                    warn!(
+                        "Unknown manual texture view for screenshot, skipping: {:?}",
+                        texture_view
+                    );
+                    continue;
+                };
+                let format = manual_texture_view.format;
+                let size = Extent3d {
+                    width: manual_texture_view.size.x,
+                    height: manual_texture_view.size.y,
+                    ..default()
+                };
+                let (texture_view, state) = prepare_screenshot_state(
+                    size,
+                    format,
+                    &render_device,
+                    &screenshot_pipeline,
+                    &pipeline_cache,
+                    &mut pipelines,
+                );
+                prepared.insert(*entity, state);
+                view_target_attachments.insert(
+                    target.clone(),
+                    OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()),
+                );
+            }
+        }
+    }
+}
+
+fn prepare_screenshot_state(
+    size: Extent3d,
+    format: TextureFormat,
+    render_device: &RenderDevice,
+    pipeline: &ScreenshotToScreenPipeline,
+    pipeline_cache: &PipelineCache,
+    pipelines: &mut SpecializedRenderPipelines<ScreenshotToScreenPipeline>,
+) -> (TextureView, ScreenshotPreparedState) {
+    let texture = render_device.create_texture(&wgpu::TextureDescriptor {
+        label: Some("screenshot-capture-rendertarget"),
+        size,
+        mip_level_count: 1,
+        sample_count: 1,
+        dimension: wgpu::TextureDimension::D2,
+        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: get_aligned_size(size.width, size.height, 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(
+        "screenshot-to-screen-bind-group",
+        &pipeline.bind_group_layout,
+        &BindGroupEntries::single(&texture_view),
+    );
+    let pipeline_id = pipelines.specialize(pipeline_cache, pipeline, format);
+
+    (
+        texture_view,
+        ScreenshotPreparedState {
+            texture,
+            buffer,
+            bind_group,
+            pipeline_id,
+            size,
+        },
+    )
+}
 
 pub struct ScreenshotPlugin;
 
@@ -128,7 +403,15 @@ const SCREENSHOT_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(11918575
 
 impl Plugin for ScreenshotPlugin {
     fn build(&self, app: &mut bevy_app::App) {
-        app.init_resource::<ScreenshotManager>();
+        app.add_systems(
+            First,
+            clear_screenshots
+                .after(event_update_system)
+                .before(apply_deferred),
+        )
+        .add_systems(Update, trigger_screenshots)
+        .register_type::<Screenshot>()
+        .register_type::<ScreenshotCaptured>();
 
         load_internal_asset!(
             app,
@@ -139,8 +422,23 @@ impl Plugin for ScreenshotPlugin {
     }
 
     fn finish(&self, app: &mut bevy_app::App) {
+        let (tx, rx) = std::sync::mpsc::channel();
+        app.insert_resource(CapturedScreenshots(Arc::new(Mutex::new(rx))));
+
         if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
-            render_app.init_resource::<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>();
+            render_app
+                .insert_resource(RenderScreenshotsSender(tx))
+                .init_resource::<RenderScreenshotTargets>()
+                .init_resource::<RenderScreenshotsPrepared>()
+                .init_resource::<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>()
+                .add_systems(ExtractSchedule, extract_screenshots.ambiguous_with_all())
+                .add_systems(
+                    Render,
+                    prepare_screenshots
+                        .after(prepare_view_attachments)
+                        .before(prepare_view_targets)
+                        .in_set(RenderSet::ManageViews),
+                );
         }
     }
 }
@@ -221,27 +519,104 @@ impl SpecializedRenderPipeline for ScreenshotToScreenPipeline {
     }
 }
 
-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 targets = world.resource::<RenderScreenshotTargets>();
+    let prepared = world.resource::<RenderScreenshotsPrepared>();
     let pipelines = world.resource::<PipelineCache>();
+    let gpu_images = world.resource::<RenderAssets<GpuImage>>();
+    let windows = world.resource::<ExtractedWindows>();
+    let manual_texture_views = world.resource::<ManualTextureViews>();
 
-    for window in windows.values() {
+    for (entity, render_target) in targets.iter() {
-        if let Some(memory) = &window.screenshot_memory {
+        match render_target {
+            NormalizedRenderTarget::Window(window) => {
+                let window = window.entity();
+                let Some(window) = windows.get(&window) else {
+                    continue;
+                };
                 let width = window.physical_width;
                 let height = window.physical_height;
-            let texture_format = window.swap_chain_texture_format.unwrap();
+                let Some(texture_format) = window.swap_chain_texture_format else {
+                    continue;
+                };
+                let Some(swap_chain_texture) = window.swap_chain_texture.as_ref() else {
+                    continue;
+                };
+                let texture_view = swap_chain_texture.texture.create_view(&Default::default());
+                render_screenshot(
+                    encoder,
+                    prepared,
+                    pipelines,
+                    entity,
+                    width,
+                    height,
+                    texture_format,
+                    &texture_view,
+                );
+            }
+            NormalizedRenderTarget::Image(image) => {
+                let Some(gpu_image) = gpu_images.get(image) else {
+                    warn!("Unknown image for screenshot, skipping: {:?}", image);
+                    continue;
+                };
+                let width = gpu_image.size.x;
+                let height = gpu_image.size.y;
+                let texture_format = gpu_image.texture_format;
+                let texture_view = gpu_image.texture_view.deref();
+                render_screenshot(
+                    encoder,
+                    prepared,
+                    pipelines,
+                    entity,
+                    width,
+                    height,
+                    texture_format,
+                    texture_view,
+                );
+            }
+            NormalizedRenderTarget::TextureView(texture_view) => {
+                let Some(texture_view) = manual_texture_views.get(texture_view) else {
+                    warn!(
+                        "Unknown manual texture view for screenshot, skipping: {:?}",
+                        texture_view
+                    );
+                    continue;
+                };
+                let width = texture_view.size.x;
+                let height = texture_view.size.y;
+                let texture_format = texture_view.format;
+                let texture_view = texture_view.texture_view.deref();
+                render_screenshot(
+                    encoder,
+                    prepared,
+                    pipelines,
+                    entity,
+                    width,
+                    height,
+                    texture_format,
+                    texture_view,
+                );
+            }
+        };
+    }
+}
 
+#[allow(clippy::too_many_arguments)]
+fn render_screenshot(
+    encoder: &mut CommandEncoder,
+    prepared: &RenderScreenshotsPrepared,
+    pipelines: &PipelineCache,
+    entity: &Entity,
+    width: u32,
+    height: u32,
+    texture_format: TextureFormat,
+    texture_view: &wgpu::TextureView,
+) {
+    if let Some(prepared_state) = &prepared.get(entity) {
         encoder.copy_texture_to_buffer(
-                memory.texture.as_image_copy(),
+            prepared_state.texture.as_image_copy(),
             wgpu::ImageCopyBuffer {
-                    buffer: &memory.buffer,
+                buffer: &prepared_state.buffer,
                 layout: layout_data(width, height, texture_format),
             },
             Extent3d {
@@ -250,17 +625,12 @@ pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEnc
                 ..Default::default()
             },
         );
-            if let Some(pipeline) = pipelines.get_render_pipeline(memory.pipeline_id) {
+
-                let true_swapchain_texture_view = window
+        if let Some(pipeline) = pipelines.get_render_pipeline(prepared_state.pipeline_id) {
-                    .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,
+                    view: texture_view,
                     resolve_target: None,
                     ops: wgpu::Operations {
                         load: wgpu::LoadOp::Load,
@@ -272,24 +642,27 @@ pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEnc
                 occlusion_query_set: None,
             });
             pass.set_pipeline(pipeline);
-                pass.set_bind_group(0, &memory.bind_group, &[]);
+            pass.set_bind_group(0, &prepared_state.bind_group, &[]);
             pass.draw(0..3, 0..1);
         }
     }
 }
-}
 
 pub(crate) fn collect_screenshots(world: &mut World) {
-    let _span = info_span!("collect_screenshots");
+    #[cfg(feature = "trace")]
+    let _span = bevy_utils::tracing::info_span!("collect_screenshots").entered();
 
-    let mut windows = world.resource_mut::<ExtractedWindows>();
+    let sender = world.resource::<RenderScreenshotsSender>().deref().clone();
-    for window in windows.values_mut() {
+    let prepared = world.resource::<RenderScreenshotsPrepared>();
-        if let Some(screenshot_func) = window.screenshot_func.take() {
+
-            let width = window.physical_width;
+    for (entity, prepared) in prepared.iter() {
-            let height = window.physical_height;
+        let entity = *entity;
-            let texture_format = window.swap_chain_texture_format.unwrap();
+        let sender = sender.clone();
+        let width = prepared.size.width;
+        let height = prepared.size.height;
+        let texture_format = prepared.texture.format();
         let pixel_size = texture_format.pixel_size();
-            let ScreenshotPreparedState { buffer, .. } = window.screenshot_memory.take().unwrap();
+        let buffer = prepared.buffer.clone();
 
         let finish = async move {
             let (tx, rx) = async_channel::bounded(1);
@@ -307,7 +680,6 @@ pub(crate) fn collect_screenshots(world: &mut World) {
             // 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.
@@ -318,17 +690,16 @@ pub(crate) fn collect_screenshots(world: &mut World) {
                 let mut take_offset = buffered_row_bytes;
                 let mut place_offset = initial_row_bytes;
                 for _ in 1..height {
-                        result.copy_within(
+                    result.copy_within(take_offset..take_offset + buffered_row_bytes, place_offset);
-                            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(
+            if let Err(e) = sender.send((
+                entity,
+                Image::new(
                     Extent3d {
                         width,
                         height,
@@ -338,10 +709,12 @@ pub(crate) fn collect_screenshots(world: &mut World) {
                     result,
                     texture_format,
                     RenderAssetUsages::RENDER_WORLD,
-                ));
+                ),
+            )) {
+                error!("Failed to send screenshot: {:?}", e);
+            }
         };
 
         AsyncComputeTaskPool::get().spawn(finish).detach();
     }
 }
-}

examples/window/screenshot.rs

@@ -1,29 +1,48 @@
 //! An example showing how to save screenshots to disk
 
 use bevy::prelude::*;
-use bevy::render::view::screenshot::ScreenshotManager;
+use bevy::window::SystemCursorIcon;
-use bevy::window::PrimaryWindow;
+use bevy_render::view::cursor::CursorIcon;
+use bevy_render::view::screenshot::{save_to_disk, Capturing, Screenshot};
 
 fn main() {
     App::new()
         .add_plugins(DefaultPlugins)
         .add_systems(Startup, setup)
-        .add_systems(Update, screenshot_on_spacebar)
+        .add_systems(Update, (screenshot_on_spacebar, screenshot_saving))
         .run();
 }
 
 fn screenshot_on_spacebar(
+    mut commands: Commands,
     input: Res<ButtonInput<KeyCode>>,
-    main_window: Query<Entity, With<PrimaryWindow>>,
-    mut screenshot_manager: ResMut<ScreenshotManager>,
     mut counter: Local<u32>,
 ) {
     if input.just_pressed(KeyCode::Space) {
         let path = format!("./screenshot-{}.png", *counter);
         *counter += 1;
-        screenshot_manager
+        commands
-            .save_screenshot_to_disk(main_window.single(), path)
+            .spawn(Screenshot::primary_window())
-            .unwrap();
+            .observe(save_to_disk(path));
+    }
+}
+
+fn screenshot_saving(
+    mut commands: Commands,
+    screenshot_saving: Query<Entity, With<Capturing>>,
+    windows: Query<Entity, With<Window>>,
+) {
+    let window = windows.single();
+    match screenshot_saving.iter().count() {
+        0 => {
+            commands.entity(window).remove::<CursorIcon>();
+        }
+        x if x > 0 => {
+            commands
+                .entity(window)
+                .insert(CursorIcon::from(SystemCursorIcon::Progress));
+        }
+        _ => {}
     }
 }