remove extract stage to run during extract

Author: hymm

crates/bevy_app/src/app.rs

@@ -189,6 +189,14 @@ impl App {
         if app.is_building_plugin {
             panic!("App::run() was called from within Plugin::Build(), which is not allowed.");
         }
+
+        // temporarily remove the plugin registry to run each plugin's setup function on app.
+        let mut plugin_registry = std::mem::take(&mut app.plugin_registry);
+        for plugin in &plugin_registry {
+            plugin.setup(&mut app);
+        }
+        std::mem::swap(&mut app.plugin_registry, &mut plugin_registry);
+
         let runner = std::mem::replace(&mut app.runner, Box::new(run_once));
         (runner)(app);
     }

crates/bevy_app/src/plugin.rs

@@ -16,10 +16,19 @@ use std::any::Any;
 pub trait Plugin: Downcast + Any + Send + Sync {
     /// Configures the [`App`] to which this plugin is added.
     fn build(&self, app: &mut App);
+
+    /// Runs after all plugins are built, but before the app runner is called.
+    /// This can be useful if you have some resource that other plugins need during their build step,
+    /// but after build you want to remove it and send it to another thread.
+    fn setup(&self, _app: &mut App) {
+        // do nothing
+    }
+
     /// Configures a name for the [`Plugin`] which is primarily used for debugging.
     fn name(&self) -> &str {
         std::any::type_name::<Self>()
     }
+
     /// If the plugin can be meaningfully instantiated several times in an [`App`](crate::App),
     /// override this method to return `false`.
     fn is_unique(&self) -> bool {

crates/bevy_ecs/src/schedule/mod.rs

@@ -361,6 +361,17 @@ impl Schedule {
             .and_then(|stage| stage.downcast_mut::<T>())
     }
 
+    /// Remove a [`Stage`] from the schedule
+    pub fn remove_stage(&mut self, stage_label: impl StageLabel) -> Option<Box<dyn Stage>> {
+        let label = stage_label.as_label();
+
+        let Some(index) = self.stage_order.iter().position(|x| *x == label) else {
+                return None;
+            };
+        self.stage_order.remove(index);
+        self.stages.remove(&label)
+    }
+
     /// Executes each [`Stage`] contained in the schedule, one at a time.
     pub fn run_once(&mut self, world: &mut World) {
         for label in &self.stage_order {

crates/bevy_render/src/lib.rs

@@ -92,6 +92,10 @@ pub enum RenderStage {
     Cleanup,
 }
 
+/// Resource for holding the extract stage of the rendering schedule
+#[derive(Resource)]
+pub struct ExtractStage(pub SystemStage);
+
 /// The simulation [`World`] of the application, stored as a resource.
 /// This resource is only available during [`RenderStage::Extract`] and not
 /// during command application of that stage.
@@ -266,6 +270,20 @@ impl Plugin for RenderPlugin {
             .register_type::<primitives::CubemapFrusta>()
             .register_type::<primitives::Frustum>();
     }
+
+    fn setup(&self, app: &mut App) {
+        if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
+            // move the extract stage to a resource so render_app.run() does not run it.
+            let stage = render_app
+                .schedule
+                .remove_stage(RenderStage::Extract)
+                .unwrap()
+                .downcast::<SystemStage>()
+                .unwrap();
+
+            render_app.world.insert_resource(ExtractStage(*stage));
+        }
+    }
 }
 
 /// A "scratch" world used to avoid allocating new worlds every frame when
@@ -276,25 +294,23 @@ struct ScratchMainWorld(World);
 /// Executes the [`Extract`](RenderStage::Extract) stage of the renderer.
 /// This updates the render world with the extracted ECS data of the current frame.
 fn extract(app_world: &mut World, render_app: &mut App) {
-    let extract = render_app
-        .schedule
-        .get_stage_mut::<SystemStage>(RenderStage::Extract)
-        .unwrap();
-
-    // temporarily add the app world to the render world as a resource
-    let scratch_world = app_world.remove_resource::<ScratchMainWorld>().unwrap();
-    let inserted_world = std::mem::replace(app_world, scratch_world.0);
-    let running_world = &mut render_app.world;
-    running_world.insert_resource(MainWorld(inserted_world));
-
-    extract.run(running_world);
-    // move the app world back, as if nothing happened.
-    let inserted_world = running_world.remove_resource::<MainWorld>().unwrap();
-    let scratch_world = std::mem::replace(app_world, inserted_world.0);
-    app_world.insert_resource(ScratchMainWorld(scratch_world));
-
-    // Note: We apply buffers (read, Commands) after the `MainWorld` has been removed from the render app's world
-    // so that in future, pipelining will be able to do this too without any code relying on it.
-    // see <https://github.com/bevyengine/bevy/issues/5082>
-    extract.apply_buffers(running_world);
+    render_app
+        .world
+        .resource_scope(|render_world, mut extract_stage: Mut<ExtractStage>| {
+            // temporarily add the app world to the render world as a resource
+            let scratch_world = app_world.remove_resource::<ScratchMainWorld>().unwrap();
+            let inserted_world = std::mem::replace(app_world, scratch_world.0);
+            render_world.insert_resource(MainWorld(inserted_world));
+
+            extract_stage.0.run(render_world);
+            // move the app world back, as if nothing happened.
+            let inserted_world = render_world.remove_resource::<MainWorld>().unwrap();
+            let scratch_world = std::mem::replace(app_world, inserted_world.0);
+            app_world.insert_resource(ScratchMainWorld(scratch_world));
+
+            // Note: We apply buffers (read, Commands) after the `MainWorld` has been removed from the render app's world
+            // so that in future, pipelining will be able to do this too without any code relying on it.
+            // see <https://github.com/bevyengine/bevy/issues/5082>
+            extract_stage.0.apply_buffers(render_world);
+        });
 }