diff --git a/crates/bevy_app/src/app.rs b/crates/bevy_app/src/app.rs index b78fbedd1a241..065babcca9d70 100644 --- a/crates/bevy_app/src/app.rs +++ b/crates/bevy_app/src/app.rs @@ -130,6 +130,10 @@ impl Default for App { .in_set(bevy_ecs::message::MessageUpdateSystems) .run_if(bevy_ecs::message::message_update_condition), ); + app.add_systems( + crate::Last, + bevy_ecs::system::despawn_unused_registered_systems, + ); app.add_message::(); app @@ -382,6 +386,24 @@ impl App { self.main_mut().register_system(system) } + /// Registers a system and returns a tracked [`SystemHandle`] so it can later + /// be called by [`World::run_system`]. The system entity will be automatically + /// queued for despawn when the last clone of the returned handle is dropped. + /// + /// See [`World::register_tracked_system`] for more details. + /// + /// [`SystemHandle`]: bevy_ecs::system::SystemHandle + pub fn register_tracked_system( + &mut self, + system: impl IntoSystem + 'static, + ) -> bevy_ecs::system::SystemHandle + where + I: SystemInput + 'static, + O: 'static, + { + self.main_mut().register_tracked_system(system) + } + /// Configures a collection of system sets in the provided schedule, adding any sets that do not exist. #[track_caller] pub fn configure_sets( @@ -2086,4 +2108,21 @@ mod tests { assert_eq!(test_events.len(), 2); // Events are double-buffered, so we see 2 + 0 = 2 assert_eq!(test_events.iter_current_update_messages().count(), 0); } + + #[test] + fn auto_despawn_unused_registered_systems() { + let mut app = App::new(); + + fn my_system() {} + + let handle = app.register_tracked_system(my_system); + let entity = handle.entity(); + + app.update(); + assert!(app.world().get_entity(entity).is_ok()); + + drop(handle); + app.update(); + assert!(app.world().get_entity(entity).is_err()); + } } diff --git a/crates/bevy_app/src/sub_app.rs b/crates/bevy_app/src/sub_app.rs index 1b20aab2ff5d3..8e52a8236ed76 100644 --- a/crates/bevy_app/src/sub_app.rs +++ b/crates/bevy_app/src/sub_app.rs @@ -247,6 +247,18 @@ impl SubApp { self.world.register_system(system) } + /// See [`App::register_tracked_system`]. + pub fn register_tracked_system( + &mut self, + system: impl IntoSystem + 'static, + ) -> bevy_ecs::system::SystemHandle + where + I: SystemInput + 'static, + O: 'static, + { + self.world.register_tracked_system(system) + } + /// See [`App::configure_sets`]. #[track_caller] pub fn configure_sets( diff --git a/crates/bevy_ecs/src/system/system_registry.rs b/crates/bevy_ecs/src/system/system_registry.rs index 819f78c8c9c44..62f4246f13225 100644 --- a/crates/bevy_ecs/src/system/system_registry.rs +++ b/crates/bevy_ecs/src/system/system_registry.rs @@ -5,13 +5,16 @@ use crate::{ entity::Entity, error::BevyError, system::{ - input::SystemInput, BoxedSystem, IntoSystem, RunSystemError, SystemParamValidationError, + input::SystemInput, BoxedSystem, Commands, If, IntoSystem, Res, RunSystemError, + SystemParamValidationError, }, world::World, }; use alloc::boxed::Box; use bevy_ecs_macros::{Component, Resource}; +use bevy_platform::sync::Arc; use bevy_utils::prelude::DebugName; +use concurrent_queue::ConcurrentQueue; use core::{any::TypeId, marker::PhantomData}; use thiserror::Error; @@ -94,6 +97,135 @@ impl RemovedSystem { } } +/// A system that despawns any registered system entities whose [`SystemHandle`] +/// reference count has reached zero. +pub fn despawn_unused_registered_systems( + // `RegisteredSystemDespawner` is initialized lazily the first time a system + // is registered, so it's possible that it doesn't exist yet when this system runs. + despawner: If>, + mut commands: Commands, +) { + for entity in despawner.queue.try_iter() { + // In case the entity was already despawned manually, we ignore the error here. + commands.entity(entity).try_despawn(); + } +} + +/// A resource that stores the channel for despawning unused registered system +/// entities. +#[derive(Resource)] +pub struct RegisteredSystemDespawner { + queue: Arc>, +} + +impl Default for RegisteredSystemDespawner { + fn default() -> Self { + Self { + queue: Arc::new(ConcurrentQueue::unbounded()), + } + } +} + +/// A maybe-strong handle to an entity acting as a registered system. Strong +/// handles are created by [`World::register_tracked_system`] or +/// [`World::register_tracked_boxed_system`]. +/// +/// Strong handles provide automatic cleanup of registered systems once all clones +/// of the handle are dropped, while weak handles do not. However, the **existence +/// of a strong handle does not prevent the registered system entity from being +/// despawned manually**, like with [`World::unregister_system`] or +/// [`World::unregister_system_cached`]. +/// +/// # Cleanup +/// +/// Registered system entities are cleaned up by the [`despawn_unused_registered_systems`] +/// system, which is automatically added to the default app by the `bevy_app` +/// crate when the "std" feature is enabled. If not using the default app, the +/// "std" feature, or `bevy_app` in general, consider running this system +/// yourself to ensure proper cleanup of registered systems. +pub enum SystemHandle { + /// A strong handle keeps the system entity alive as long as the handle + /// (and any clones of it) exist, as long as the system entity isn't + /// manually despawned. + Strong(Arc), + /// A weak handle does not keep the system entity alive. + Weak(SystemId), +} + +impl SystemHandle { + /// Returns the [`Entity`] of the registered system associated with this handle. + pub fn entity(&self) -> Entity { + match self { + SystemHandle::Strong(strong) => strong.entity, + SystemHandle::Weak(weak) => weak.entity, + } + } +} + +impl Eq for SystemHandle {} + +// A manual impl is used because the trait bounds should ignore the `I` and `O` phantom parameters. +impl Clone for SystemHandle { + fn clone(&self) -> Self { + match self { + SystemHandle::Strong(strong) => SystemHandle::Strong(Arc::clone(strong)), + SystemHandle::Weak(weak) => SystemHandle::Weak(*weak), + } + } +} + +// A manual impl is used because the trait bounds should ignore the `I` and `O` phantom parameters, +// and so that strong and weak handles can be compared for equality based on their entities. +impl PartialEq for SystemHandle { + fn eq(&self, other: &Self) -> bool { + self.entity() == other.entity() + } +} + +impl PartialEq> for SystemHandle { + fn eq(&self, other: &SystemId) -> bool { + self.entity() == other.entity + } +} + +// A manual impl is used because the trait bounds should ignore the `I` and `O` phantom parameters, +// and so that the handle can be hashed based on its entity not its handle type. +impl core::hash::Hash for SystemHandle { + fn hash(&self, state: &mut H) { + self.entity().hash(state); + } +} + +impl core::fmt::Debug for SystemHandle { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + let name = if matches!(self, SystemHandle::Strong(_)) { + "StrongSystemHandle" + } else { + "WeakSystemHandle" + }; + f.debug_tuple(name).field(&self.entity()).finish() + } +} + +impl From> for SystemHandle { + fn from(id: SystemId) -> Self { + SystemHandle::Weak(id) + } +} + +/// A strong handle for a registered system that despawns the entity when dropped. +pub struct StrongSystemHandle { + entity: Entity, + drop_queue: Arc>, +} + +impl Drop for StrongSystemHandle { + fn drop(&mut self) { + // Send the entity to be despawned by the world when the last strong handle is dropped. + let _ = self.drop_queue.push(self.entity); + } +} + /// An identifier for a registered system. /// /// These are opaque identifiers, keyed to a specific [`World`], @@ -214,6 +346,51 @@ impl World { SystemId::from_entity(entity) } + /// Registers a system and returns a tracked [`SystemHandle`] so it can later + /// be called by [`World::run_system`]. The system entity will be automatically + /// queued for despawn when the last clone of the returned handle is dropped. + /// + /// By default, unused tracked system entities are despawned by the + /// [`despawn_unused_registered_systems`] system in the `Last` schedule of + /// the default app. Otherwise, it needs to be run manually to ensure proper + /// cleanup of registered systems. + /// + /// It's possible to register multiple copies of the same system by calling + /// this function multiple times. If that's not what you want, consider using + /// [`World::register_system_cached`] instead. + pub fn register_tracked_system( + &mut self, + system: impl IntoSystem + 'static, + ) -> SystemHandle + where + I: SystemInput + 'static, + O: 'static, + { + self.register_tracked_boxed_system(Box::new(IntoSystem::into_system(system))) + } + + /// Similar to [`Self::register_tracked_system`], but allows passing in a + /// [`BoxedSystem`]. + /// + /// This is useful if the [`IntoSystem`] implementor has already been turned + /// into a [`System`](crate::system::System) trait object and put in a [`Box`]. + pub fn register_tracked_boxed_system( + &mut self, + system: BoxedSystem, + ) -> SystemHandle + where + I: SystemInput + 'static, + O: 'static, + { + let entity = self.spawn(RegisteredSystem::new(system)).id(); + let despawner = self.get_resource_or_init::(); + + SystemHandle::Strong(Arc::new(StrongSystemHandle { + entity, + drop_queue: despawner.queue.clone(), + })) + } + /// Removes a registered system and returns the system, if it exists. /// After removing a system, the [`SystemId`] becomes invalid and attempting to use it afterwards will result in errors. /// Re-adding the removed system will register it on a new [`SystemId`]. @@ -332,7 +509,7 @@ impl World { /// ``` pub fn run_system( &mut self, - id: SystemId<(), O>, + id: impl Into>, ) -> Result> { self.run_system_with(id, ()) } @@ -364,13 +541,14 @@ impl World { /// See [`World::run_system`] for more examples. pub fn run_system_with( &mut self, - id: SystemId, + id: impl Into>, input: I::Inner<'_>, ) -> Result> where I: SystemInput + 'static, O: 'static, { + let id = id.into(); // Lookup let mut entity = self .get_entity_mut(id.entity) @@ -604,7 +782,7 @@ mod tests { use crate::{ prelude::*, - system::{RegisteredSystemError, SystemId}, + system::{despawn_unused_registered_systems, RegisteredSystemError, SystemId}, }; #[derive(Resource, Default, PartialEq, Debug)] @@ -1074,4 +1252,23 @@ mod tests { Err(err) => panic!("Failed with wrong error. `{:?}`", err), } } + + #[test] + fn despawn_unused() { + let mut world = World::new(); + + fn system() {} + + let handle = world.register_tracked_system(system); + let entity = handle.entity(); + drop(handle); + + assert!(world.get_entity(entity).is_ok()); + + world + .run_system_cached(despawn_unused_registered_systems) + .unwrap(); + + assert!(world.get_entity(entity).is_err()); + } }