Skip to content

Commit

Permalink
Remove InstanceId when Scene Despawn (#12778)
Browse files Browse the repository at this point in the history
# 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<DynamicScene>`, and when the
`Handle<DynamicScene>` component is removed, delete the corresponding
`InstanceId` from `spawned_dynamic_scenes`
  • Loading branch information
Fpgu authored and mockersf committed Apr 1, 2024
1 parent 0341944 commit ae7e40c
Show file tree
Hide file tree
Showing 3 changed files with 83 additions and 15 deletions.
2 changes: 1 addition & 1 deletion crates/bevy_scene/src/bundle.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
///
Expand Down
33 changes: 32 additions & 1 deletion crates/bevy_scene/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,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)]
Expand All @@ -52,6 +52,37 @@ impl Plugin for ScenePlugin {
.add_event::<SceneInstanceReady>()
.init_resource::<SceneSpawner>()
.add_systems(SpawnScene, (scene_spawner, scene_spawner_system).chain());

// Register component hooks for DynamicScene
app.world
.register_component_hooks::<Handle<DynamicScene>>()
.on_remove(|mut world, entity, _| {
let Some(handle) = world.get::<Handle<DynamicScene>>(entity) else {
return;
};
let id = handle.id();
if let Some(&SceneInstance(scene_instance)) = world.get::<SceneInstance>(entity) {
let Some(mut scene_spawner) = world.get_resource_mut::<SceneSpawner>() 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::<Handle<Scene>>()
.on_remove(|mut world, entity, _| {
if let Some(&SceneInstance(scene_instance)) = world.get::<SceneInstance>(entity) {
let Some(mut scene_spawner) = world.get_resource_mut::<SceneSpawner>() else {
return;
};
scene_spawner.despawn_instance(scene_instance);
}
});
}
}

Expand Down
63 changes: 50 additions & 13 deletions crates/bevy_scene/src/scene_spawner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,9 +60,8 @@ impl InstanceId {
/// - [`despawn_instance`](Self::despawn_instance)
#[derive(Default, Resource)]
pub struct SceneSpawner {
spawned_scenes: HashMap<AssetId<Scene>, Vec<InstanceId>>,
spawned_dynamic_scenes: HashMap<AssetId<DynamicScene>, Vec<InstanceId>>,
spawned_instances: HashMap<InstanceId, InstanceInfo>,
pub(crate) spawned_dynamic_scenes: HashMap<AssetId<DynamicScene>, HashSet<InstanceId>>,
pub(crate) spawned_instances: HashMap<InstanceId, InstanceInfo>,
scene_asset_event_reader: ManualEventReader<AssetEvent<DynamicScene>>,
dynamic_scenes_to_spawn: Vec<(Handle<DynamicScene>, InstanceId)>,
scenes_to_spawn: Vec<(Handle<Scene>, InstanceId)>,
Expand Down Expand Up @@ -207,7 +206,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)
}

Expand Down Expand Up @@ -248,8 +247,6 @@ impl SceneSpawner {
scene.write_to_world_with(world, &world.resource::<AppTypeRegistry>().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)
})
}
Expand Down Expand Up @@ -307,8 +304,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));
Expand Down Expand Up @@ -440,17 +437,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::*;

Expand Down Expand Up @@ -551,4 +545,47 @@ mod tests {
},
);
}

#[test]
fn despawn_scene() {
let mut app = App::new();
app.add_plugins((AssetPlugin::default(), ScenePlugin));
app.register_type::<ComponentA>();

let asset_server = app.world.resource::<AssetServer>();

// 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::<SceneSpawner>();
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<Entity, With<ComponentA>>| {
for entity in query.iter() {
commands.entity(entity).despawn_recursive();
}
},
);

app.update();
check(&mut app.world, 0);
}
}

0 comments on commit ae7e40c

Please sign in to comment.