From 3d0f2409d5fefce70a7e6346ddca5e88d7d52284 Mon Sep 17 00:00:00 2001
From: Josh Robson Chase <josh@robsonchase.com>
Date: Mon, 23 Sep 2024 20:08:10 -0500
Subject: [PATCH] bevy_ecs: flush entities after running observers and hooks in
 despawn (#15398)

# Objective

Fixes #14467

Observers and component lifecycle hooks are allowed to perform
operations that subsequently require `Entities` to be flushed, such as
reserving a new entity. If this occurs during an `on_remove` hook or an
`OnRemove` event trigger during an `EntityWorldMut::despawn`, a panic
will occur.

## Solution

Call `world.flush_entities()` after running `on_remove` hooks/observers
during `despawn`

## Testing

Added a new test that fails before the fix and succeeds afterward.
---
 crates/bevy_ecs/src/observer/mod.rs     | 22 ++++++++++++++++++++++
 crates/bevy_ecs/src/world/entity_ref.rs |  5 ++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/crates/bevy_ecs/src/observer/mod.rs b/crates/bevy_ecs/src/observer/mod.rs
index a339bf6b1f479..3b87f59291e27 100644
--- a/crates/bevy_ecs/src/observer/mod.rs
+++ b/crates/bevy_ecs/src/observer/mod.rs
@@ -1160,4 +1160,26 @@ mod tests {
         world.flush();
         assert_eq!(vec!["event", "event"], world.resource::<Order>().0);
     }
+
+    // Regression test for https://github.com/bevyengine/bevy/issues/14467
+    // Fails prior to https://github.com/bevyengine/bevy/pull/15398
+    #[test]
+    fn observer_on_remove_during_despawn_spawn_empty() {
+        let mut world = World::new();
+
+        // Observe the removal of A - this will run during despawn
+        world.observe(|_: Trigger<OnRemove, A>, mut cmd: Commands| {
+            // Spawn a new entity - this reserves a new ID and requires a flush
+            // afterward before Entities::free can be called.
+            cmd.spawn_empty();
+        });
+
+        let ent = world.spawn(A).id();
+
+        // Despawn our entity, which runs the OnRemove observer and allocates a
+        // new Entity.
+        // Should not panic - if it does, then Entities was not flushed properly
+        // after the observer's spawn_empty.
+        world.despawn(ent);
+    }
 }
diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs
index c91885469c766..6fd1fa5b2518c 100644
--- a/crates/bevy_ecs/src/world/entity_ref.rs
+++ b/crates/bevy_ecs/src/world/entity_ref.rs
@@ -1297,7 +1297,6 @@ impl<'w> EntityWorldMut<'w> {
     /// See [`World::despawn`] for more details.
     pub fn despawn(self) {
         let world = self.world;
-        world.flush_entities();
         let archetype = &world.archetypes[self.location.archetype_id];
 
         // SAFETY: Archetype cannot be mutably aliased by DeferredWorld
@@ -1323,6 +1322,10 @@ impl<'w> EntityWorldMut<'w> {
             world.removed_components.send(component_id, self.entity);
         }
 
+        // Observers and on_remove hooks may reserve new entities, which
+        // requires a flush before Entities::free may be called.
+        world.flush_entities();
+
         let location = world
             .entities
             .free(self.entity)