Schedule enter and exit hooks

crates/bevy_ecs/src/schedule/hook.rs

@@ -0,0 +1,174 @@
+//! This is about hooks for the [`Schedule`](super::Schedule) execution phases,
+//! aiming to handle instructions triggered either before entering the [`Schedule`] or after exiting it.
+use crate::world::World;
+use crate::{intern::Interned, prelude::Resource, schedule::ScheduleLabel, system::SystemId};
+use bevy_platform::collections::HashMap;
+use bevy_platform::prelude::vec::Vec;
+use log::error;
+
+/// Used to control whether to retain or remove the [`ScheduleHook`] after it is triggered.
+#[derive(Debug, Clone, PartialEq, Eq, Copy, Hash)]
+pub enum ScheduleHookPlan {
+    /// Remove after executing the [`ScheduleHook`]
+    Clear,
+    /// Keep after executing the [`ScheduleHook`]
+    Keep,
+}
+
+/// Every valid [`Schedule`](super::Schedule) hook is a system that returns a [ScheduleHookPlan].
+pub type ScheduleHook = SystemId<(), ScheduleHookPlan>;
+
+/// The hub for managing [`ScheduleHook`], used to control when hooks are triggered.
+#[derive(Debug, Resource, Default, Clone)]
+pub struct ScheduleHooks {
+    enter: HashMap<Interned<dyn ScheduleLabel>, Vec<ScheduleHook>>,
+    exit: HashMap<Interned<dyn ScheduleLabel>, Vec<ScheduleHook>>,
+}
+
+impl ScheduleHooks {
+    /// Add a [`ScheduleHook`] to a [`ScheduleLabel`] that triggers before entering the [`Schedule`](super::Schedule).
+    pub fn add_enter_hook(&mut self, label: impl ScheduleLabel, hook: ScheduleHook) -> &mut Self {
+        self.enter
+            .entry(label.intern())
+            .and_modify(|hooks| {
+                hooks.push(hook);
+            })
+            .or_insert(Vec::from([hook]));
+        self
+    }
+
+    /// Add a [`ScheduleHook`] to a [`ScheduleLabel`] that triggers after exiting the [`Schedule`](super::Schedule).
+    pub fn add_exit_hook(&mut self, label: impl ScheduleLabel, hook: ScheduleHook) -> &mut Self {
+        self.exit
+            .entry(label.intern())
+            .and_modify(|hooks| {
+                hooks.push(hook);
+            })
+            .or_insert(Vec::from([hook]));
+        self
+    }
+
+    /// Execute the [`ScheduleHook`] that runs before a [`ScheduleLabel`].
+    pub fn run_enter(&mut self, world: &mut World, label: impl ScheduleLabel) {
+        if let Some(hooks) = self.enter.get_mut(&label.intern()) {
+            hooks.retain(|hook| {
+                world
+                    .run_system(hook.clone())
+                    .unwrap_or_else(|err| {
+                        error!(
+                            "a enter schedule hook fail,schedule label: {:?}, system id:{:?},error context:{:?}",
+                            label, hook, err
+                        );
+                        ScheduleHookPlan::Clear
+                    })
+                    .eq(&ScheduleHookPlan::Keep)
+            });
+        }
+    }
+
+    /// Execute the [`ScheduleHook`] that runs after a [`ScheduleLabel`].
+    pub fn run_exit(&mut self, world: &mut World, label: impl ScheduleLabel) {
+        if let Some(hooks) = self.exit.get_mut(&label.intern()) {
+            hooks.retain(|hook| {
+                world
+                    .run_system(hook.clone())
+                    .unwrap_or_else(|err| {
+                        error!(
+                            "a exit schedule hook fail,schedule label: {:?}, system id:{:?},error context:{:?}",
+                            label, hook, err
+                        );
+                        ScheduleHookPlan::Clear
+                    })
+                    .eq(&ScheduleHookPlan::Keep)
+            });
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        prelude::Component,
+        system::{Commands, Local},
+    };
+
+    use super::*;
+
+    #[derive(Debug, ScheduleLabel, Hash, Clone, PartialEq, Eq)]
+    pub struct HookLabel;
+
+    #[derive(Debug, Component, PartialEq, Eq)]
+    pub struct TestComponent;
+
+    pub const SPAWN_COUNT: usize = 4;
+
+    #[test]
+    fn hook_success_run() {
+        let mut world = World::new();
+
+        let system = world.register_system(|mut commands: Commands, mut count: Local<usize>| {
+            commands.spawn(TestComponent);
+            if *count < SPAWN_COUNT {
+                *count += 1;
+                ScheduleHookPlan::Keep
+            } else {
+                ScheduleHookPlan::Clear
+            }
+        });
+
+        let mut hooks = ScheduleHooks::default();
+
+        hooks.add_enter_hook(HookLabel, system);
+
+        for _ in 0..SPAWN_COUNT {
+            hooks.run_enter(&mut world, HookLabel);
+        }
+
+        let mut query = world.query::<&TestComponent>();
+
+        let iter = query.iter(&world);
+
+        assert_eq!(SPAWN_COUNT, iter.count());
+
+        hooks.run_enter(&mut world, HookLabel);
+
+        assert!(hooks
+            .enter
+            .get(&HookLabel.intern())
+            .is_some_and(|hooks| hooks.is_empty()));
+    }
+
+    #[test]
+    fn hook_fail_run() {
+        let mut world = World::new();
+
+        let system = world.register_system(|mut commands: Commands| {
+            commands.spawn(TestComponent);
+            ScheduleHookPlan::Clear
+        });
+
+        let mut hooks = ScheduleHooks::default();
+
+        hooks.add_enter_hook(HookLabel, system);
+
+        assert_eq!(
+            Some(1),
+            hooks
+                .enter
+                .get(&HookLabel.intern())
+                .map(|hooks| hooks.len())
+        );
+
+        world.despawn(system.entity);
+
+        hooks.run_enter(&mut world, HookLabel);
+
+        assert_eq!(
+            Some(0),
+            hooks
+                .enter
+                .get(&HookLabel.intern())
+                .map(|hooks| hooks.len())
+        );
+    }
+}

crates/bevy_ecs/src/schedule/mod.rs

@@ -5,14 +5,17 @@ mod condition;
 mod config;
 mod error;
 mod executor;
+mod hook;
 mod node;
 mod pass;
 mod schedule;
 mod set;
 mod stepping;
 
 pub use self::graph::GraphInfo;
-pub use self::{condition::*, config::*, error::*, executor::*, node::*, schedule::*, set::*};
+pub use self::{
+    condition::*, config::*, error::*, executor::*, hook::*, node::*, schedule::*, set::*,
+};
 pub use pass::{FlattenedDependencies, ScheduleBuildPass};
 
 /// An implementation of a graph data structure.

crates/bevy_ecs/src/world/mod.rs

@@ -58,7 +58,7 @@ use crate::{
     query::{DebugCheckedUnwrap, QueryData, QueryFilter, QueryState},
     relationship::RelationshipHookMode,
     resource::{IsResource, Resource, ResourceEntities, IS_RESOURCE},
-    schedule::{Schedule, ScheduleLabel, Schedules},
+    schedule::{Schedule, ScheduleHooks, ScheduleLabel, Schedules},
     storage::{NonSendData, Storages},
     system::Commands,
     world::{
@@ -3729,6 +3729,11 @@ impl World {
         schedules.insert(schedule);
     }
 
+    /// Get [`ScheduleHooks`], otherwise initialize and get.
+    pub fn schedule_hooks(&mut self) -> Mut<'_, ScheduleHooks> {
+        self.get_resource_or_init::<ScheduleHooks>()
+    }
+
     /// Temporarily removes the schedule associated with `label` from the world,
     /// runs user code, and finally re-adds the schedule.
     /// This returns a [`TryRunScheduleError`] if there is no schedule
@@ -3820,7 +3825,17 @@ impl World {
         &mut self,
         label: impl ScheduleLabel,
     ) -> Result<(), TryRunScheduleError> {
-        self.try_schedule_scope(label, |world, sched| sched.run(world))
+        self.try_schedule_scope(label.intern(), |world, sched| {
+            world.try_resource_scope::<ScheduleHooks, ()>(|world, mut hooks| {
+                hooks.run_enter(world, label.intern());
+            });
+
+            sched.run(world);
+
+            world.try_resource_scope::<ScheduleHooks, ()>(|world, mut hooks| {
+                hooks.run_exit(world, label);
+            });
+        })
     }
 
     /// Runs the [`Schedule`] associated with the `label` a single time.
@@ -3835,7 +3850,17 @@ impl World {
     ///
     /// If the requested schedule does not exist.
     pub fn run_schedule(&mut self, label: impl ScheduleLabel) {
-        self.schedule_scope(label, |world, sched| sched.run(world));
+        self.schedule_scope(label.intern(), |world, sched| {
+            world.try_resource_scope::<ScheduleHooks, ()>(|world, mut hooks| {
+                hooks.run_enter(world, label.intern());
+            });
+
+            sched.run(world);
+
+            world.try_resource_scope::<ScheduleHooks, ()>(|world, mut hooks| {
+                hooks.run_exit(world, label);
+            });
+        });
     }
 
     /// Ignore system order ambiguities caused by conflicts on [`Component`]s of type `T`.
@@ -3927,11 +3952,14 @@ mod tests {
     use crate::{
         change_detection::{DetectChangesMut, MaybeLocation},
         component::{ComponentCloneBehavior, ComponentDescriptor, ComponentInfo, StorageType},
-        entity::EntityHashSet,
+        entity::{Entity, EntityHashSet},
         entity_disabling::{DefaultQueryFilters, Disabled},
-        prelude::{Event, Mut, On, Res},
+        prelude::{Event, Mut, On, Res, Schedule},
         ptr::OwningPtr,
+        query::{Changed, With},
         resource::Resource,
+        schedule::ScheduleHookPlan,
+        system::{Commands, Query},
         world::{error::EntityMutableFetchError, DeferredWorld},
     };
     use alloc::{
@@ -3941,7 +3969,7 @@ mod tests {
         vec,
         vec::Vec,
     };
-    use bevy_ecs_macros::Component;
+    use bevy_ecs_macros::{Component, ScheduleLabel};
     use bevy_platform::collections::{HashMap, HashSet};
     use bevy_utils::prelude::DebugName;
     use core::{
@@ -4633,4 +4661,60 @@ mod tests {
 
         assert!(world.get_entity(eid).is_err());
     }
+
+    #[test]
+    fn schedules_hook() {
+        #[derive(Component, PartialEq, Debug)]
+        struct Foo;
+
+        #[derive(Component, PartialEq, Debug)]
+        struct Bar;
+
+        #[derive(ScheduleLabel, PartialEq, Debug, Eq, Clone, Hash)]
+        struct HookLabel;
+
+        pub const SPAWN_COUNT: usize = 3;
+
+        let mut world = World::new();
+
+        let mut schedule = Schedule::new(HookLabel);
+
+        schedule.add_systems(|query: Query<(), With<Bar>>| {
+            assert_eq!(SPAWN_COUNT, query.count());
+        });
+
+        world.add_schedule(schedule);
+
+        let enter_hook =
+            world.register_system(|mut commands: Commands, query: Query<(), Changed<Foo>>| {
+                let count = query.count();
+                commands.spawn_batch((0..count).map(|_| Bar));
+                ScheduleHookPlan::Clear
+            });
+
+        let exit_hook =
+            world.register_system(|mut commands: Commands, query: Query<Entity, With<Foo>>| {
+                for foo in query {
+                    commands.entity(foo).despawn();
+                }
+                ScheduleHookPlan::Clear
+            });
+
+        world
+            .schedule_hooks()
+            .add_enter_hook(HookLabel, enter_hook)
+            .add_exit_hook(HookLabel, exit_hook);
+
+        world.spawn_batch((0..SPAWN_COUNT).map(|_| Foo));
+
+        world.run_schedule(HookLabel);
+
+        let mut foo_query = world.query::<&Foo>();
+
+        assert_eq!(0, foo_query.iter(&world).count());
+
+        let mut bar_query = world.query::<&Bar>();
+
+        assert_eq!(SPAWN_COUNT, bar_query.iter(&world).count());
+    }
 }