Exit schedule for systems to run at end of app run

crates/bevy_app/src/app.rs

@@ -1,4 +1,6 @@
-use crate::{CoreStage, Plugin, PluginGroup, StartupSchedule, StartupStage};
+use crate::{
+    CoreStage, ExitSchedule, ExitStage, Plugin, PluginGroup, StartupSchedule, StartupStage,
+};
 pub use bevy_derive::AppLabel;
 use bevy_ecs::{
     event::{Event, Events},
@@ -221,6 +223,16 @@ impl App {
         self.world.clear_trackers();
     }
 
+    /// Executes the [`App`]'s [exit schedule](Self::add_default_stages()). Should be called by the runner once right before it completes.
+    pub fn teardown(&mut self) {
+        #[cfg(feature = "trace")]
+        let _bevy_frame_update_span = info_span!("exit app").entered();
+        if let Some(schedule) = self.schedule.get_stage_mut::<Schedule>(ExitSchedule) {
+            schedule.set_run_criteria(ShouldRun::once);
+            schedule.run(&mut self.world);
+        }
+    }
+
     /// Starts the application by calling the app's [runner function](Self::set_runner).
     ///
     /// Finalizes the [`App`] configuration. For general usage, see the example on the item
@@ -675,6 +687,62 @@ impl App {
         self
     }
 
+    /// Adds a system to the [exit schedule](Self::add_default_stages)
+    ///
+    /// If the code you want to run only requires access to a single resource/component, consider
+    /// if implementing the [`Drop`] trait on the component or resource might work. This can be more
+    /// idiomatic for objects representing system resources.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use bevy_app::prelude::*;
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// # let mut app = App::new();
+    /// # fn my_exit_system() {}
+    /// #
+    /// app.add_exit_system(my_exit_system);
+    /// ```
+    pub fn add_exit_system<Params>(
+        &mut self,
+        system: impl IntoSystemDescriptor<Params>,
+    ) -> &mut Self {
+        self.schedule
+            .stage(ExitSchedule, |schedule: &mut Schedule| {
+                schedule.add_system_to_stage(ExitStage, system)
+            });
+        self
+    }
+
+    /// Adds a [`SystemSet`] to the [exit schedule](Self::add_default_stages)
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use bevy_app::prelude::*;
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// # let mut app = App::new();
+    /// # fn exit_system_a() {}
+    /// # fn exit_system_b() {}
+    /// # fn exit_system_c() {}
+    /// #
+    /// app.add_exit_system_set(
+    ///     SystemSet::new()
+    ///         .with_system(exit_system_a)
+    ///         .with_system(exit_system_b)
+    ///         .with_system(exit_system_c),
+    /// );
+    /// ```
+    pub fn add_exit_system_set(&mut self, system_set: SystemSet) -> &mut Self {
+        self.schedule
+            .stage(ExitSchedule, |schedule: &mut Schedule| {
+                schedule.add_system_set_to_stage(ExitStage, system_set)
+            });
+        self
+    }
+
     /// Adds a new [`State`] with the given `initial` value.
     /// This inserts a new `State<T>` resource and adds a new "driver" to [`CoreStage::Update`].
     /// Each stage that uses `State<T>` for system run criteria needs a driver. If you need to use
@@ -725,6 +793,8 @@ impl App {
     /// - **Post-update:** Often used by plugins to finalize their internal state after the
     ///   world changes that happened during the update stage.
     /// - **Last:** Runs right before the end of the schedule execution cycle.
+    /// - **Exit:** This schedule doesn't execute during the normal schedule execution cycle,
+    ///   but is instead executed right before the game exits.
     ///
     /// The labels for those stages are defined in the [`CoreStage`] and [`StartupStage`] `enum`s.
     ///
@@ -749,6 +819,12 @@ impl App {
             .add_stage(CoreStage::Update, SystemStage::parallel())
             .add_stage(CoreStage::PostUpdate, SystemStage::parallel())
             .add_stage(CoreStage::Last, SystemStage::parallel())
+            .add_stage(
+                ExitSchedule,
+                Schedule::default()
+                    .with_run_criteria(|| ShouldRun::No)
+                    .with_stage(ExitStage, SystemStage::parallel()),
+            )
     }
 
     /// Setup the application to manage events of type `T`.
@@ -1166,7 +1242,11 @@ fn run_once(mut app: App) {
 ///
 /// You can also use this event to detect that an exit was requested. In order to receive it, systems
 /// subscribing to this event should run after it was emitted and before the schedule of the same
-/// frame is over. This is important since [`App::run()`] might never return.
+/// frame is over. This is important since [`App::run()`] might close the program before returning.
+///
+/// It would probably be easier just to add an [exit system](`App::add_exit_system()`), which is
+/// added to [a special schedule](`App::add_default_stages()`) for runners to execute right before
+/// the App closes.
 ///
 /// If you don't require access to other components or resources, consider implementing the [`Drop`]
 /// trait on components/resources for code that runs on exit. That saves you from worrying about

crates/bevy_app/src/lib.rs

@@ -64,3 +64,12 @@ pub enum StartupStage {
     /// The [`Stage`](bevy_ecs::schedule::Stage) that runs once after [`StartupStage::Startup`].
     PostStartup,
 }
+
+/// The label for the exit [`Schedule`](bevy_ecs::schedule::Schedule),
+/// which runs once at the end of the [`App`].
+#[derive(Debug, Hash, PartialEq, Eq, Clone, StageLabel)]
+pub struct ExitSchedule;
+
+/// The only stage in the exit schedule.
+#[derive(Debug, Hash, PartialEq, Eq, Clone, StageLabel)]
+pub struct ExitStage;

crates/bevy_app/src/schedule_runner.rs

@@ -95,9 +95,11 @@ impl Plugin for ScheduleRunnerPlugin {
                         if let Some(app_exit_events) =
                             app.world.get_resource_mut::<Events<AppExit>>()
                         {
-                            if let Some(exit) = app_exit_event_reader.iter(&app_exit_events).last()
+                            if let Some(exit) =
+                                app_exit_event_reader.iter(&app_exit_events).last().cloned()
                             {
-                                return Err(exit.clone());
+                                app.teardown();
+                                return Err(exit);
                             }
                         }
 
@@ -106,9 +108,11 @@ impl Plugin for ScheduleRunnerPlugin {
                         if let Some(app_exit_events) =
                             app.world.get_resource_mut::<Events<AppExit>>()
                         {
-                            if let Some(exit) = app_exit_event_reader.iter(&app_exit_events).last()
+                            if let Some(exit) =
+                                app_exit_event_reader.iter(&app_exit_events).last().cloned()
                             {
-                                return Err(exit.clone());
+                                app.teardown();
+                                return Err(exit);
                             }
                         }
 
@@ -168,3 +172,50 @@ impl Plugin for ScheduleRunnerPlugin {
         });
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{App, AppExit, ScheduleRunnerPlugin};
+    use bevy_ecs::{
+        prelude::EventWriter,
+        system::{Res, ResMut, Resource},
+    };
+
+    #[derive(Resource)]
+    struct RunCounter(usize);
+
+    #[test]
+    fn exit_systems_called_after_app_exits() {
+        use std::sync::{
+            atomic::{AtomicBool, Ordering},
+            Arc,
+        };
+
+        let successful_exit = Arc::new(AtomicBool::new(false));
+        let exit_marker = Arc::clone(&successful_exit);
+
+        App::new()
+            .insert_resource(RunCounter(0))
+            .add_plugin(ScheduleRunnerPlugin)
+            .add_system(
+                |mut exit: EventWriter<AppExit>, mut run_counter: ResMut<RunCounter>| {
+                    run_counter.0 += 1;
+                    if run_counter.0 == 3 {
+                        exit.send(AppExit);
+                    }
+                },
+            )
+            .add_exit_system(move |run_counter: Res<RunCounter>| {
+                if run_counter.0 != 3 {
+                    panic!("Exit systems should only run at the end of the program");
+                }
+                exit_marker.store(true, Ordering::SeqCst);
+            })
+            .run();
+
+        assert!(
+            successful_exit.load(Ordering::SeqCst),
+            "Exit system did not run successfully"
+        );
+    }
+}

crates/bevy_winit/src/lib.rs

@@ -260,6 +260,7 @@ pub fn winit_runner(mut app: App) {
 
         if let Some(app_exit_events) = app.world.get_resource::<Events<AppExit>>() {
             if app_exit_event_reader.iter(app_exit_events).last().is_some() {
+                app.teardown();
                 *control_flow = ControlFlow::Exit;
                 return;
             }