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 <[email protected]>

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_utils::tracing::{debug, info, warn};
-use bevy_window::PrimaryWindow;
+use bevy_render::view::screenshot::{save_to_disk, Screenshot};
+use bevy_utils::tracing::{debug, info};
 
 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
-                    .save_screenshot_to_disk(main_window, path)
-                    .unwrap();
+                world
+                    .spawn(Screenshot::primary_window())
+                    .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)]
-pub fn prepare_view_targets(
-    mut commands: Commands,
+/// Prepares the view target [`OutputColorAttachment`] for each view in the current frame.
+pub fn prepare_view_attachments(
     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(|| {
-            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 {
+        let Some(out_attachment) = view_target_attachments.get(target) 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::{
-        BindGroupEntries, PipelineCache, SpecializedRenderPipelines, SurfaceTexture, TextureView,
-    },
+    render_resource::{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,
-    TextureViewDescriptor,
+    SurfaceConfiguration, SurfaceTargetUnsafe, TextureFormat, TextureUsages, 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,
-            });
-        }
     }
 }