From c227fc9fada2a327d174644a3ce7e5c210aecfb6 Mon Sep 17 00:00:00 2001
From: Brian Reavis <brian@thirdroute.com>
Date: Fri, 26 Jan 2024 19:16:57 -0800
Subject: [PATCH] Fix infinite asset preparation due to undrained AssetEvent
 events  (#11383)

# Objective

After #10520, I was experiencing seriously degraded performance that
ended up being due to never-drained `AssetEvent` events causing havoc
inside `extract_render_asset::<A>`. The same events being read over and
over again meant the same assets were being prepared every frame for
eternity. For what it's worth, I was noticing this on a static scene
about every 3rd or so time running my project.

* References #10520
* Fixes #11240

Why these events aren't sometimes drained between frames is beyond me
and perhaps worthy of another investigation, but the approach in this PR
effectively restores the original cached `EventReader` behavior (which
fixes it).

## Solution

I followed the [`CachedSystemState`
example](https://github.com/bevyengine/bevy/blob/3a666cab23eab5e83552ae074f47d8663221311e/crates/bevy_ecs/src/system/function_system.rs#L155)
to make sure that the `EventReader` state is cached between frames like
it used to be when it was an argument of `extract_render_asset::<A>`.
---
 crates/bevy_render/src/render_asset.rs | 92 ++++++++++++++++----------
 1 file changed, 57 insertions(+), 35 deletions(-)

diff --git a/crates/bevy_render/src/render_asset.rs b/crates/bevy_render/src/render_asset.rs
index 2ea2caaf5e..2226e849b2 100644
--- a/crates/bevy_render/src/render_asset.rs
+++ b/crates/bevy_render/src/render_asset.rs
@@ -5,6 +5,7 @@ use bevy_ecs::{
     prelude::{Commands, EventReader, IntoSystemConfigs, ResMut, Resource},
     schedule::SystemConfigs,
     system::{StaticSystemParam, SystemParam, SystemParamItem, SystemState},
+    world::{FromWorld, Mut},
 };
 use bevy_reflect::Reflect;
 use bevy_utils::{thiserror::Error, HashMap, HashSet};
@@ -88,6 +89,7 @@ impl<A: RenderAsset, AFTER: RenderAssetDependency + 'static> Plugin
     for RenderAssetPlugin<A, AFTER>
 {
     fn build(&self, app: &mut App) {
+        app.init_resource::<CachedExtractRenderAssetSystemState<A>>();
         if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
             render_app
                 .init_resource::<ExtractedAssets<A>>()
@@ -176,49 +178,69 @@ impl<A: RenderAsset> RenderAssets<A> {
     }
 }
 
+#[derive(Resource)]
+struct CachedExtractRenderAssetSystemState<A: RenderAsset> {
+    state: SystemState<(
+        EventReader<'static, 'static, AssetEvent<A>>,
+        ResMut<'static, Assets<A>>,
+    )>,
+}
+
+impl<A: RenderAsset> FromWorld for CachedExtractRenderAssetSystemState<A> {
+    fn from_world(world: &mut bevy_ecs::world::World) -> Self {
+        Self {
+            state: SystemState::new(world),
+        }
+    }
+}
+
 /// This system extracts all created or modified assets of the corresponding [`RenderAsset`] type
 /// into the "render world".
 fn extract_render_asset<A: RenderAsset>(mut commands: Commands, mut main_world: ResMut<MainWorld>) {
-    let mut system_state: SystemState<(EventReader<AssetEvent<A>>, ResMut<Assets<A>>)> =
-        SystemState::new(&mut main_world);
-    let (mut events, mut assets) = system_state.get_mut(&mut main_world);
+    main_world.resource_scope(
+        |world, mut cached_state: Mut<CachedExtractRenderAssetSystemState<A>>| {
+            let (mut events, mut assets) = cached_state.state.get_mut(world);
 
-    let mut changed_assets = HashSet::default();
-    let mut removed = Vec::new();
-    for event in events.read() {
-        #[allow(clippy::match_same_arms)]
-        match event {
-            AssetEvent::Added { id } | AssetEvent::Modified { id } => {
-                changed_assets.insert(*id);
-            }
-            AssetEvent::Removed { .. } => {}
-            AssetEvent::Unused { id } => {
-                changed_assets.remove(id);
-                removed.push(*id);
-            }
-            AssetEvent::LoadedWithDependencies { .. } => {
-                // TODO: handle this
-            }
-        }
-    }
+            let mut changed_assets = HashSet::default();
+            let mut removed = Vec::new();
 
-    let mut extracted_assets = Vec::new();
-    for id in changed_assets.drain() {
-        if let Some(asset) = assets.get(id) {
-            if asset.persistence_policy() == RenderAssetPersistencePolicy::Unload {
-                if let Some(asset) = assets.remove(id) {
-                    extracted_assets.push((id, asset));
+            for event in events.read() {
+                #[allow(clippy::match_same_arms)]
+                match event {
+                    AssetEvent::Added { id } | AssetEvent::Modified { id } => {
+                        changed_assets.insert(*id);
+                    }
+                    AssetEvent::Removed { .. } => {}
+                    AssetEvent::Unused { id } => {
+                        changed_assets.remove(id);
+                        removed.push(*id);
+                    }
+                    AssetEvent::LoadedWithDependencies { .. } => {
+                        // TODO: handle this
+                    }
                 }
-            } else {
-                extracted_assets.push((id, asset.clone()));
             }
-        }
-    }
 
-    commands.insert_resource(ExtractedAssets {
-        extracted: extracted_assets,
-        removed,
-    });
+            let mut extracted_assets = Vec::new();
+            for id in changed_assets.drain() {
+                if let Some(asset) = assets.get(id) {
+                    if asset.persistence_policy() == RenderAssetPersistencePolicy::Unload {
+                        if let Some(asset) = assets.remove(id) {
+                            extracted_assets.push((id, asset));
+                        }
+                    } else {
+                        extracted_assets.push((id, asset.clone()));
+                    }
+                }
+            }
+
+            commands.insert_resource(ExtractedAssets {
+                extracted: extracted_assets,
+                removed,
+            });
+            cached_state.state.apply(world);
+        },
+    );
 }
 
 // TODO: consider storing inside system?