Tracked one-shot systems

crates/bevy_app/src/app.rs

@@ -130,6 +130,11 @@ impl Default for App {
                 .in_set(bevy_ecs::message::MessageUpdateSystems)
                 .run_if(bevy_ecs::message::message_update_condition),
         );
+        #[cfg(feature = "std")]
+        app.add_systems(
+            crate::Last,
+            bevy_ecs::system::despawn_unused_registered_systems,
+        );
         app.add_message::<AppExit>();
 
         app
@@ -382,6 +387,25 @@ 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
+    #[cfg(feature = "std")]
+    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 +2110,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,19 @@ impl SubApp {
         self.world.register_system(system)
     }
 
+    /// See [`App::register_tracked_system`].
+    #[cfg(feature = "std")]
+    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/Cargo.toml

@@ -69,6 +69,7 @@ std = [
   "bevy_utils/std",
   "bitflags/std",
   "concurrent-queue/std",
+  "crossbeam-channel/std",
   "fixedbitset/std",
   "indexmap/std",
   "serde?/std",
@@ -101,6 +102,7 @@ bevy_platform = { path = "../bevy_platform", version = "0.19.0-dev", default-fea
 ] }
 
 bitflags = { version = "2.3", default-features = false }
+crossbeam-channel = { version = "0.5", default-features = false }
 fixedbitset = { version = "0.5", default-features = false }
 serde = { version = "1", default-features = false, features = [
   "alloc",

crates/bevy_ecs/src/system/system_registry.rs

@@ -11,6 +11,8 @@ use crate::{
 };
 use alloc::boxed::Box;
 use bevy_ecs_macros::{Component, Resource};
+#[cfg(feature = "std")]
+use bevy_platform::sync::Arc;
 use bevy_utils::prelude::DebugName;
 use core::{any::TypeId, marker::PhantomData};
 use thiserror::Error;
@@ -94,6 +96,150 @@ impl<I, O> RemovedSystem<I, O> {
     }
 }
 
+/// A system that despawns any registered system entities whose [`SystemHandle`]
+/// reference count has reached zero.
+#[cfg(feature = "std")]
+pub fn despawn_unused_registered_systems(
+    despawner: Option<crate::system::Res<RegisteredSystemDespawner>>,
+    mut commands: crate::system::Commands,
+) {
+    // `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.
+    if let Some(despawner) = despawner {
+        for entity in despawner.receiver.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.
+#[cfg(feature = "std")]
+#[derive(Resource)]
+pub struct RegisteredSystemDespawner {
+    receiver: crossbeam_channel::Receiver<Entity>,
+    sender: crossbeam_channel::Sender<Entity>,
+}
+
+#[cfg(feature = "std")]
+impl Default for RegisteredSystemDespawner {
+    fn default() -> Self {
+        let (sender, receiver) = crossbeam_channel::unbounded();
+        Self { receiver, sender }
+    }
+}
+
+/// 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.
+#[cfg(feature = "std")]
+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>),
+}
+
+#[cfg(feature = "std")]
+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,
+        }
+    }
+}
+
+#[cfg(feature = "std")]
+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.
+#[cfg(feature = "std")]
+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.
+#[cfg(feature = "std")]
+impl<I: SystemInput, O> PartialEq for SystemHandle<I, O> {
+    fn eq(&self, other: &Self) -> bool {
+        self.entity() == other.entity()
+    }
+}
+
+#[cfg(feature = "std")]
+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.
+#[cfg(feature = "std")]
+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);
+    }
+}
+
+#[cfg(feature = "std")]
+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()
+    }
+}
+
+#[cfg(feature = "std")]
+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 keeps the system entity alive
+/// as long as the handle (and any clones of it) exist.
+#[cfg(feature = "std")]
+pub struct StrongSystemHandle {
+    entity: Entity,
+    drop_sender: crossbeam_channel::Sender<Entity>,
+}
+
+#[cfg(feature = "std")]
+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_sender.send(self.entity);
+    }
+}
+
 /// An identifier for a registered system.
 ///
 /// These are opaque identifiers, keyed to a specific [`World`],
@@ -214,6 +360,55 @@ 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.
+    #[cfg(feature = "std")]
+    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`].
+    #[cfg(feature = "std")]
+    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();
+        #[cfg(feature = "std")]
+        let despawner = self.get_resource_or_init::<RegisteredSystemDespawner>();
+
+        SystemHandle::Strong(Arc::new(StrongSystemHandle {
+            entity,
+            #[cfg(feature = "std")]
+            drop_sender: despawner.sender.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 +527,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 +559,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 +800,7 @@ mod tests {
 
     use crate::{
         prelude::*,
-        system::{RegisteredSystemError, SystemId},
+        system::{despawn_unused_registered_systems, RegisteredSystemError, SystemId},
     };
 
     #[derive(Resource, Default, PartialEq, Debug)]
@@ -1074,4 +1270,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());
+
+        let mut cleanup = IntoSystem::into_system(despawn_unused_registered_systems);
+        cleanup.initialize(&mut world);
+        cleanup.run((), &mut world).unwrap();
+
+        assert!(world.get_entity(entity).is_err());
+    }
 }