Separate Extract from Sub App Schedule (#7046)

# Objective

- This pulls out some of the changes to Plugin setup and sub apps from #6503 to make that PR easier to review.
- Separate the extract stage from running the sub app's schedule to allow for them to be run on separate threads in the future
- Fixes #6990

## Solution

- add a run method to `SubApp` that runs the schedule
- change the name of `sub_app_runner` to extract to make it clear that this function is only for extracting data between the main app and the sub app
- remove the extract stage from the sub app schedule so it can be run separately. This is done by adding a `setup` method to the `Plugin` trait that runs after all plugin build methods run. This is required to allow the extract stage to be removed from the schedule after all the plugins have added their systems to the stage. We will also need the setup method for pipelined rendering to setup the render thread. See https://github.com/bevyengine/bevy/blob/e3267965e15f14be18eec942dcaf16807144eb05/crates/bevy_render/src/pipelined_rendering.rs#L57-L98

## Changelog

- Separate SubApp Extract stage from running the sub app schedule.

## Migration Guide

### SubApp `runner` has conceptually been changed to an `extract` function.

The `runner` no longer is in charge of running the sub app schedule. It's only concern is now moving data between the main world and the sub app. The `sub_app.app.schedule` is now run for you after the provided function is called.

```rust
// before
fn main() {
    let sub_app = App::empty();
    sub_app.add_stage(MyStage, SystemStage::parallel());
    
    App::new().add_sub_app(MySubApp, sub_app, move |main_world, sub_app| {
        extract(app_world, render_app);
        render_app.app.schedule.run();
    });
}

// after
fn main() {
        let sub_app = App::empty();
    sub_app.add_stage(MyStage, SystemStage::parallel());
    
    App::new().add_sub_app(MySubApp, sub_app, move |main_world, sub_app| {
        extract(app_world, render_app);
        // schedule is automatically called for you after extract is run
    });
}
```

Author: hymm

crates/bevy_app/src/app.rs

@@ -90,7 +90,20 @@ impl Debug for App {
 /// Each `SubApp` has its own [`Schedule`] and [`World`], enabling a separation of concerns.
 struct SubApp {
     app: App,
-    runner: Box<dyn Fn(&mut World, &mut App)>,
+    extract: Box<dyn Fn(&mut World, &mut App)>,
+}
+
+impl SubApp {
+    /// Runs the `SubApp`'s schedule.
+    pub fn run(&mut self) {
+        self.app.schedule.run(&mut self.app.world);
+        self.app.world.clear_trackers();
+    }
+
+    /// Extracts data from main world to this sub-app.
+    pub fn extract(&mut self, main_world: &mut World) {
+        (self.extract)(main_world, &mut self.app);
+    }
 }
 
 impl Debug for SubApp {
@@ -153,8 +166,8 @@ impl App {
         self.schedule.run(&mut self.world);
 
         for sub_app in self.sub_apps.values_mut() {
-            (sub_app.runner)(&mut self.world, &mut sub_app.app);
-            sub_app.app.world.clear_trackers();
+            sub_app.extract(&mut self.world);
+            sub_app.run();
         }
 
         self.world.clear_trackers();
@@ -176,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);
     }
@@ -1004,13 +1025,13 @@ impl App {
         &mut self,
         label: impl AppLabel,
         app: App,
-        sub_app_runner: impl Fn(&mut World, &mut App) + 'static,
+        extract: impl Fn(&mut World, &mut App) + 'static,
     ) -> &mut Self {
         self.sub_apps.insert(
             label.as_label(),
             SubApp {
                 app,
-                runner: Box::new(sub_app_runner),
+                extract: Box::new(extract),
             },
         );
         self

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>())
     }
 
+    /// Removes 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.
@@ -198,7 +202,10 @@ impl Plugin for RenderPlugin {
                         .with_system(PipelineCache::process_pipeline_queue_system)
                         .with_system(render_system.at_end()),
                 )
-                .add_stage(RenderStage::Cleanup, SystemStage::parallel())
+                .add_stage(
+                    RenderStage::Cleanup,
+                    SystemStage::parallel().with_system(World::clear_entities.at_end()),
+                )
                 .init_resource::<render_graph::RenderGraph>()
                 .insert_resource(RenderInstance(instance))
                 .insert_resource(device)
@@ -248,78 +255,6 @@ impl Plugin for RenderPlugin {
                     // extract
                     extract(app_world, render_app);
                 }
-
-                {
-                    #[cfg(feature = "trace")]
-                    let _stage_span =
-                        bevy_utils::tracing::info_span!("stage", name = "prepare").entered();
-
-                    // prepare
-                    let prepare = render_app
-                        .schedule
-                        .get_stage_mut::<SystemStage>(RenderStage::Prepare)
-                        .unwrap();
-                    prepare.run(&mut render_app.world);
-                }
-
-                {
-                    #[cfg(feature = "trace")]
-                    let _stage_span =
-                        bevy_utils::tracing::info_span!("stage", name = "queue").entered();
-
-                    // queue
-                    let queue = render_app
-                        .schedule
-                        .get_stage_mut::<SystemStage>(RenderStage::Queue)
-                        .unwrap();
-                    queue.run(&mut render_app.world);
-                }
-
-                {
-                    #[cfg(feature = "trace")]
-                    let _stage_span =
-                        bevy_utils::tracing::info_span!("stage", name = "sort").entered();
-
-                    // phase sort
-                    let phase_sort = render_app
-                        .schedule
-                        .get_stage_mut::<SystemStage>(RenderStage::PhaseSort)
-                        .unwrap();
-                    phase_sort.run(&mut render_app.world);
-                }
-
-                {
-                    #[cfg(feature = "trace")]
-                    let _stage_span =
-                        bevy_utils::tracing::info_span!("stage", name = "render").entered();
-
-                    // render
-                    let render = render_app
-                        .schedule
-                        .get_stage_mut::<SystemStage>(RenderStage::Render)
-                        .unwrap();
-                    render.run(&mut render_app.world);
-                }
-
-                {
-                    #[cfg(feature = "trace")]
-                    let _stage_span =
-                        bevy_utils::tracing::info_span!("stage", name = "cleanup").entered();
-
-                    // cleanup
-                    let cleanup = render_app
-                        .schedule
-                        .get_stage_mut::<SystemStage>(RenderStage::Cleanup)
-                        .unwrap();
-                    cleanup.run(&mut render_app.world);
-                }
-                {
-                    #[cfg(feature = "trace")]
-                    let _stage_span =
-                        bevy_utils::tracing::info_span!("stage", name = "clear_entities").entered();
-
-                    render_app.world.clear_entities();
-                }
             });
         }
 
@@ -335,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
@@ -345,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);
+        });
 }