From cdecd39e313a17a1c389e2e7f17157b3b949aa22 Mon Sep 17 00:00:00 2001 From: Fpgu <35809384+341101@users.noreply.github.com> Date: Sun, 31 Mar 2024 06:16:49 +0800 Subject: [PATCH] Remove `InstanceId` when Scene Despawn (#12778) # Objective - Fix #12746 - When users despawn a scene, the `InstanceId` within `spawned_scenes` and `spawned_dynamic_scenes` is not removed, causing a potential memory leak ## Solution - `spawned_scenes` field was never used, and I removed it - Add a component remove hook for `Handle`, and when the `Handle` component is removed, delete the corresponding `InstanceId` from `spawned_dynamic_scenes` --- crates/bevy_scene/src/bundle.rs | 2 +- crates/bevy_scene/src/lib.rs | 33 +++++++++++++- crates/bevy_scene/src/scene_spawner.rs | 63 ++++++++++++++++++++------ 3 files changed, 83 insertions(+), 15 deletions(-) diff --git a/crates/bevy_scene/src/bundle.rs b/crates/bevy_scene/src/bundle.rs index b50abbc0e1e3d..b1d4bba4f575a 100644 --- a/crates/bevy_scene/src/bundle.rs +++ b/crates/bevy_scene/src/bundle.rs @@ -16,7 +16,7 @@ use crate::{DynamicScene, InstanceId, Scene, SceneSpawner}; /// [`InstanceId`] of a spawned scene. It can be used with the [`SceneSpawner`] to /// interact with the spawned scene. #[derive(Component, Deref, DerefMut)] -pub struct SceneInstance(InstanceId); +pub struct SceneInstance(pub(crate) InstanceId); /// A component bundle for a [`Scene`] root. /// diff --git a/crates/bevy_scene/src/lib.rs b/crates/bevy_scene/src/lib.rs index 0c1a89785412f..a05e9a5f8619e 100644 --- a/crates/bevy_scene/src/lib.rs +++ b/crates/bevy_scene/src/lib.rs @@ -44,7 +44,7 @@ pub mod prelude { } use bevy_app::prelude::*; -use bevy_asset::AssetApp; +use bevy_asset::{AssetApp, Handle}; /// Plugin that provides scene functionality to an [`App`]. #[derive(Default)] @@ -59,6 +59,37 @@ impl Plugin for ScenePlugin { .add_event::() .init_resource::() .add_systems(SpawnScene, (scene_spawner, scene_spawner_system).chain()); + + // Register component hooks for DynamicScene + app.world + .register_component_hooks::>() + .on_remove(|mut world, entity, _| { + let Some(handle) = world.get::>(entity) else { + return; + }; + let id = handle.id(); + if let Some(&SceneInstance(scene_instance)) = world.get::(entity) { + let Some(mut scene_spawner) = world.get_resource_mut::() else { + return; + }; + if let Some(instance_ids) = scene_spawner.spawned_dynamic_scenes.get_mut(&id) { + instance_ids.remove(&scene_instance); + } + scene_spawner.despawn_instance(scene_instance); + } + }); + + // Register component hooks for Scene + app.world + .register_component_hooks::>() + .on_remove(|mut world, entity, _| { + if let Some(&SceneInstance(scene_instance)) = world.get::(entity) { + let Some(mut scene_spawner) = world.get_resource_mut::() else { + return; + }; + scene_spawner.despawn_instance(scene_instance); + } + }); } } diff --git a/crates/bevy_scene/src/scene_spawner.rs b/crates/bevy_scene/src/scene_spawner.rs index e89f7f307a139..5aa918e2c0ea0 100644 --- a/crates/bevy_scene/src/scene_spawner.rs +++ b/crates/bevy_scene/src/scene_spawner.rs @@ -60,9 +60,8 @@ impl InstanceId { /// - [`despawn_instance`](Self::despawn_instance) #[derive(Default, Resource)] pub struct SceneSpawner { - spawned_scenes: HashMap, Vec>, - spawned_dynamic_scenes: HashMap, Vec>, - spawned_instances: HashMap, + pub(crate) spawned_dynamic_scenes: HashMap, HashSet>, + pub(crate) spawned_instances: HashMap, scene_asset_event_reader: ManualEventReader>, dynamic_scenes_to_spawn: Vec<(Handle, InstanceId)>, scenes_to_spawn: Vec<(Handle, InstanceId)>, @@ -210,7 +209,7 @@ impl SceneSpawner { self.spawned_instances .insert(instance_id, InstanceInfo { entity_map }); let spawned = self.spawned_dynamic_scenes.entry(id).or_default(); - spawned.push(instance_id); + spawned.insert(instance_id); Ok(instance_id) } @@ -251,8 +250,6 @@ impl SceneSpawner { scene.write_to_world_with(world, &world.resource::().clone())?; self.spawned_instances.insert(instance_id, instance_info); - let spawned = self.spawned_scenes.entry(id).or_default(); - spawned.push(instance_id); Ok(instance_id) }) } @@ -310,8 +307,8 @@ impl SceneSpawner { let spawned = self .spawned_dynamic_scenes .entry(handle.id()) - .or_insert_with(Vec::new); - spawned.push(instance_id); + .or_insert_with(HashSet::new); + spawned.insert(instance_id); } Err(SceneSpawnError::NonExistentScene { .. }) => { self.dynamic_scenes_to_spawn.push((handle, instance_id)); @@ -443,17 +440,14 @@ pub fn scene_spawner_system(world: &mut World) { mod tests { use bevy_app::App; use bevy_asset::{AssetPlugin, AssetServer}; - use bevy_ecs::component::Component; - use bevy_ecs::entity::Entity; use bevy_ecs::event::EventReader; use bevy_ecs::prelude::ReflectComponent; use bevy_ecs::query::With; - use bevy_ecs::reflect::AppTypeRegistry; use bevy_ecs::system::{Commands, Res, ResMut, RunSystemOnce}; - use bevy_ecs::world::World; + use bevy_ecs::{component::Component, system::Query}; use bevy_reflect::Reflect; - use crate::{DynamicScene, DynamicSceneBuilder, SceneInstanceReady, ScenePlugin, SceneSpawner}; + use crate::{DynamicSceneBuilder, ScenePlugin}; use super::*; @@ -554,4 +548,47 @@ mod tests { }, ); } + + #[test] + fn despawn_scene() { + let mut app = App::new(); + app.add_plugins((AssetPlugin::default(), ScenePlugin)); + app.register_type::(); + + let asset_server = app.world.resource::(); + + // Build scene. + let scene = asset_server.add(DynamicScene::default()); + let count = 10; + + // Checks the number of scene instances stored in `SceneSpawner`. + let check = |world: &mut World, expected_count: usize| { + let scene_spawner = world.resource::(); + assert_eq!( + scene_spawner.spawned_dynamic_scenes[&scene.id()].len(), + expected_count + ); + assert_eq!(scene_spawner.spawned_instances.len(), expected_count); + }; + + // Spawn scene. + for _ in 0..count { + app.world.spawn((ComponentA, scene.clone())); + } + + app.update(); + check(&mut app.world, count); + + // Despawn scene. + app.world.run_system_once( + |mut commands: Commands, query: Query>| { + for entity in query.iter() { + commands.entity(entity).despawn_recursive(); + } + }, + ); + + app.update(); + check(&mut app.world, 0); + } }