Automatic cleanup of one-shot systems

[ItsDoot]

Objective

#24087 introduces scene templating for SystemIds, however it can result in a memory leak if a scene is re-constructed multiple times:

#24087 (comment)

This was proposed basically 1:1 in #24026 (this was later changed though). The issue is that it's unclear who owns these systems, that is who is responsible for unregistering them once they are no longer needed. Given that recreating the template will spawn the system again this basically becomes a memory leak.

#24087 (comment)

Hm, are you sure this is the case even tho in build_template it only registers the system the first time its called, switching over to storing the SystemId after the first call?

If you recreate the template (e.g. you call my_scene() again) then you will create a new instance of the system. And since the system is not scoped to the scene once the scene is despawned the system entity will be leaked.

Essentially, we need a way to connect the lifetime of the registered system to the lifetime of the scene.

Solution

Introducing: SystemHandles.

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.
    Strong(Arc<StrongSystemHandle>),
    /// A weak handle does not keep the system entity alive. If a weak handle is
    /// returned by a registration function, the system entity must be
    /// manually despawned.
    Weak(SystemId<I, O>),
}

pub struct StrongSystemHandle {
    entity: Entity,
    #[cfg(feature = "std")]
    drop_sender: crossbeam_channel::Sender<Entity>,
}

Similar to bevy_asset::Handles,SystemHandle's custom Drop implementation enqueues the registered system entity into a message channel. The system despawn_unused_registered_systems pulls from the other end of this message channel and despawns the registered system entities.

World::register_system family of functions return strong handles, while Commands::register_system returns weak handles (due to the limitation of world access).

Testing

• Added a test to ensure that despawn_unused_registered_systems does its job
• Added a test to ensure that the default app will automatically call despawn_unused_registered_systems

Future work

• SystemId scene templating #24087 will use this PR as a base
• Might be worth adding a NoCleanup marker component that's attachable to registered system entities, and filter out entities with this component inside despawn_unused_registered_systems
• Probably worth adding a way for no-std folks to get automatic cleanup as well, since currently its gated on std.

[alice-i-cecile]

Strongly in favor of this core idea!

[chescock]

Why use this for cached systems? The Arcs are extra overhead, and most use cases of cached systems will never want to despawn them.

[chescock]

World::register_system and Commands::register_system returning the same type but having different behavior will be very surprising! Users might get used to the cleanup from the World version, then use Commands somewhere and introduce a leak.

I also expect we'll have cases that won't need automatic cleanup, and it would be nice to avoid the overhead of Arcs in those cases.

Could we expose separate methods on World to register automatically-cleaned and manually-cleaned systems, with different names and returning different types? Then the method on Commands could correspond directly to a method on World.

[ItsDoot]

I think switching Commands::register_system back to returning SystemId instead of SystemHandle would be the simplest fix here; users will still be able to wrap SystemIds into SystemHandles themselves, but do so explicitly. Thoughts?

[chescock]

I think switching Commands::register_system back to returning SystemId instead of SystemHandle would be the simplest fix here; users will still be able to wrap SystemIds into SystemHandles themselves, but do so explicitly. Thoughts?

That might help, although those are both impl Into<SystemId> and could be passed directly to run_system without changes.

I do also think we'll want to expose a method on World that returns a manually-cleaned system, rather than having to tell users that world.commands().register_system(s); world.flush(); is the simplest way to do it.

So I would still vote for giving them different names, even if we don't offer the manually-cleaned version on World. But I also have no authority here and am happy to be overridden :).

[chescock]

Oh, and I hope I don't sound negative on the overall idea. Automatic cleanup seems like a powerful pattern! I just think it should be decoupled from register_system a bit, since register_system is useful without automatic cleanup. (And automatic cleanup will be useful for other kinds of entities!)

[chescock]

I see how we got here, but looking at the whole thing together: Why even store the system as an entity in the World? As long as we're allocating Arcs, why not just use Arc<Mutex<dyn System>>? That gets cleaned up automatically without needing a channel or anything.

[ItsDoot]

I think there's some apprehension to requiring locking, as noted in #24072. That PR works as-is however in case we want to explore that option instead.

[chescock]

I think there's some apprehension to requiring locking, as noted in #24072. That PR works as-is however in case we want to explore that option instead.

Oh, I just meant that client code could use Arc themselves instead of using registered systems. Once you have an Arc that you can share, there's no need to store anything in the World!

... Hmm, I bet we could get rid of the locking there if we needed. The reason we don't have to lock today is that the system is owned by the World, so the &mut World gives us access to the system. But we could simulate that with a SyncUnsafeCell that has a safety requirement that it only be accessed by code with &mut World for the right world! We'd still need synchronization to set the WorldId, but that would only be on the first run.

Something like:

struct SystemArcInner<S: ?Sized> {
    world_id: OnceLock<WorldId>,
    running: SyncUnsafeCell<bool>,
    system: SyncUnsafeCell<S>,
}

impl<S: System> SystemArcInner {
    fn run(&self, world: &mut World) {
        // This takes locks the first time it's run, but is an ordinary load afterwards
        let world_id = self.world_id.get_or_init(|| {
            // This is safe because `Once` ensures only one thread runs an initializer
            unsafe { &mut *self.system.get() }.initialize(world);
            world.id()
        });
        assert_eq!(world.id(), world_id);

        // Claim the `running` lock
        // This is safe because we only access `running` when holding `&mut World` for *this* `World`
        let running = unsafe { &mut *self.running.get() };
        assert!(!*running);
        *running = true;

        // This is safe because we have the `running` lock
        let system = unsafe { &mut *self.system.get() };
        system.run(world);

        // Release the `running` lock
        // Note that this needs a new reference to `running` in case of reentrancy!
        // TODO: This should be in a `Drop` impl so it runs on unwind
        unsafe { *self.running.get() = false };
    }
}