Tracked one-shot systems

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::<AppExit>();
 
         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<I, O, M>(
+        &mut self,
+        system: impl IntoSystem<I, O, M> + 'static,
+    ) -> bevy_ecs::system::SystemHandle<I, O>
+    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<M>(
@@ -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());
+    }
 }

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<I, O, M>(
+        &mut self,
+        system: impl IntoSystem<I, O, M> + 'static,
+    ) -> bevy_ecs::system::SystemHandle<I, O>
+    where
+        I: SystemInput + 'static,
+        O: 'static,
+    {
+        self.world.register_tracked_system(system)
+    }
+
     /// See [`App::configure_sets`].
     #[track_caller]
     pub fn configure_sets<M>(

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<I, O> RemovedSystem<I, O> {
     }
 }
 
+/// 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<Res<RegisteredSystemDespawner>>,
+    mut commands: Commands,
+) {
+    while let Ok(entity) = despawner.queue.pop() {
+        // 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<ConcurrentQueue<Entity>>,
+}
+
+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<I: SystemInput = (), O = ()> {
+    /// 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<StrongSystemHandle>),
+    /// A weak handle does not keep the system entity alive.
+    Weak(SystemId<I, O>),
+}
+
+impl<I: SystemInput, O> SystemHandle<I, O> {
+    /// 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<I: SystemInput, O> Eq for SystemHandle<I, O> {}
+
+// A manual impl is used because the trait bounds should ignore the `I` and `O` phantom parameters.
+impl<I: SystemInput, O> Clone for SystemHandle<I, O> {
+    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<I: SystemInput, O> PartialEq for SystemHandle<I, O> {
+    fn eq(&self, other: &Self) -> bool {
+        self.entity() == other.entity()
+    }
+}
+
+impl<I: SystemInput, O> PartialEq<SystemId<I, O>> for SystemHandle<I, O> {
+    fn eq(&self, other: &SystemId<I, O>) -> 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<I: SystemInput, O> core::hash::Hash for SystemHandle<I, O> {
+    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
+        self.entity().hash(state);
+    }
+}
+
+impl<I: SystemInput, O> core::fmt::Debug for SystemHandle<I, O> {
+    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<I: SystemInput, O> From<SystemId<I, O>> for SystemHandle<I, O> {
+    fn from(id: SystemId<I, O>) -> 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<ConcurrentQueue<Entity>>,
+}
+
+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<I, O, M>(
+        &mut self,
+        system: impl IntoSystem<I, O, M> + 'static,
+    ) -> SystemHandle<I, O>
+    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<I, O>(
+        &mut self,
+        system: BoxedSystem<I, O>,
+    ) -> SystemHandle<I, O>
+    where
+        I: SystemInput + 'static,
+        O: 'static,
+    {
+        let entity = self.spawn(RegisteredSystem::new(system)).id();
+        let despawner = self.get_resource_or_init::<RegisteredSystemDespawner>();
+
+        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<O: 'static>(
         &mut self,
-        id: SystemId<(), O>,
+        id: impl Into<SystemId<(), O>>,
     ) -> Result<O, RegisteredSystemError<(), O>> {
         self.run_system_with(id, ())
     }
@@ -364,13 +541,14 @@ impl World {
     /// See [`World::run_system`] for more examples.
     pub fn run_system_with<I, O>(
         &mut self,
-        id: SystemId<I, O>,
+        id: impl Into<SystemId<I, O>>,
         input: I::Inner<'_>,
     ) -> Result<O, RegisteredSystemError<I, O>>
     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());
+    }
 }