Panic safe component drop

crates/bevy_ecs/src/bundle/info.rs

@@ -5,7 +5,11 @@ use bevy_platform::{
 };
 use bevy_ptr::{MovingPtr, OwningPtr};
 use bevy_utils::TypeIdMap;
-use core::{any::TypeId, ptr::NonNull};
+use core::{
+    any::{Any, TypeId},
+    panic::AssertUnwindSafe,
+    ptr::NonNull,
+};
 use indexmap::{IndexMap, IndexSet};
 
 use crate::{
@@ -212,6 +216,8 @@ impl BundleInfo {
 
     /// This writes components from a given [`Bundle`] to the given entity.
     ///
+    /// If overwritten components panic during drop, the panic payload is returned.
+    ///
     /// # Safety
     ///
     /// `bundle_component_status` must return the "correct" [`ComponentStatus`] for each component
@@ -248,7 +254,9 @@ impl BundleInfo {
         bundle: MovingPtr<'_, T>,
         insert_mode: InsertMode,
         caller: MaybeLocation,
-    ) {
+    ) -> Result<(), Box<dyn Any + Send>> {
+        let mut panic = Ok(());
+
         // NOTE: get_components calls this closure on each component in "bundle order".
         // bundle_info.component_ids are also in "bundle order"
         let mut bundle_component = 0;
@@ -258,7 +266,7 @@ impl BundleInfo {
                 .get_unchecked(bundle_component);
             // SAFETY: bundle_component is a valid index for this bundle
             let status = unsafe { bundle_component_status.get_status(bundle_component) };
-            match storage_type {
+            let f = || match storage_type {
                 StorageType::Table => {
                     let column =
                         // SAFETY: If component_id is in self.component_ids, BundleInfo::new ensures that
@@ -294,7 +302,13 @@ impl BundleInfo {
                         }
                     }
                 }
+            };
+
+            let maybe_panic = bevy_utils::catch_unwind_if_available(AssertUnwindSafe(f));
+            if panic.is_ok() {
+                panic = maybe_panic;
             }
+
             bundle_component += 1;
         });
 
@@ -308,6 +322,8 @@ impl BundleInfo {
                 caller,
             );
         }
+
+        panic
     }
 
     /// Internal method to initialize a required component from an [`OwningPtr`]. This should ultimately be called

crates/bevy_ecs/src/bundle/insert.rs

@@ -1,6 +1,7 @@
+use alloc::boxed::Box;
 use alloc::vec::Vec;
 use bevy_ptr::{ConstNonNull, MovingPtr};
-use core::ptr::NonNull;
+use core::{any::Any, ptr::NonNull};
 
 use crate::{
     archetype::{
@@ -272,6 +273,7 @@ impl<'w> BundleInserter<'w> {
                         result.table_row,
                     )
                 };
+                // No components have been dropped, so we don't need to check for a panic.
 
                 let new_location = new_archetype.allocate(entity, move_result.new_row);
                 entities.update_existing_location(entity.index(), Some(new_location));
@@ -315,6 +317,8 @@ impl<'w> BundleInserter<'w> {
         }
     }
 
+    /// Returns the entity's new location and potentially a caught panic.
+    ///
     /// # Safety
     /// - `entity` must currently exist in the source archetype for this inserter.
     /// - `location` must be `entity`'s location in the archetype.
@@ -334,10 +338,10 @@ impl<'w> BundleInserter<'w> {
         insert_mode: InsertMode,
         caller: MaybeLocation,
         relationship_hook_mode: RelationshipHookMode,
-    ) -> EntityLocation {
+    ) -> (EntityLocation, Result<(), Box<dyn Any + Send>>) {
         let archetype_after_insert = self.archetype_after_insert.as_ref();
 
-        let (new_archetype, new_location) = {
+        let (new_archetype, new_location, maybe_panic) = {
             // Non-generic prelude extracted to improve compile time by minimizing monomorphized code.
             let (new_archetype, new_location, sparse_sets, table) = Self::before_insert(
                 entity,
@@ -351,7 +355,7 @@ impl<'w> BundleInserter<'w> {
                 &mut self.archetype_move_type,
             );
 
-            self.bundle_info.as_ref().write_components(
+            let maybe_panic = self.bundle_info.as_ref().write_components(
                 table,
                 sparse_sets,
                 archetype_after_insert,
@@ -364,7 +368,7 @@ impl<'w> BundleInserter<'w> {
                 caller,
             );
 
-            (new_archetype, new_location)
+            (new_archetype, new_location, maybe_panic)
         };
 
         // SAFETY: We have no outstanding mutable references to world as they were dropped
@@ -382,7 +386,7 @@ impl<'w> BundleInserter<'w> {
             deferred_world,
         );
 
-        new_location
+        (new_location, maybe_panic)
     }
 
     // A non-generic postlude to insert used to minimize duplicated monomorphized code.

crates/bevy_ecs/src/bundle/remove.rs

@@ -1,6 +1,8 @@
+use alloc::boxed::Box;
 use alloc::vec::Vec;
 use bevy_ptr::ConstNonNull;
 use core::ptr::NonNull;
+use core::{any::Any, panic::AssertUnwindSafe};
 
 use crate::{
     archetype::{Archetype, ArchetypeCreated, ArchetypeId, Archetypes},
@@ -26,6 +28,12 @@ pub(crate) struct BundleRemover<'w> {
     pub(crate) relationship_hook_mode: RelationshipHookMode,
 }
 
+pub(crate) struct BundleRemoveResult<T> {
+    pub new_location: EntityLocation,
+    pub data: T,
+    pub panic_payload: Result<(), Box<dyn Any + Send>>,
+}
+
 impl<'w> BundleRemover<'w> {
     /// Creates a new [`BundleRemover`], if such a remover would do anything.
     ///
@@ -133,7 +141,7 @@ impl<'w> BundleRemover<'w> {
             &Components,
             &[ComponentId],
         ) -> (bool, T),
-    ) -> (EntityLocation, T) {
+    ) -> BundleRemoveResult<T> {
         // Hooks
         // SAFETY: all bundle components exist in World
         unsafe {
@@ -202,6 +210,11 @@ impl<'w> BundleRemover<'w> {
             self.bundle_info.as_ref().explicit_components(),
         );
 
+        // Component's drop functions may panic.
+        // We mustn't leave the world in an inconsistent state if that happens.
+        // Catch any such panics, finish the removal, and rethrow the first one.
+        let mut panic_payload = Ok(());
+
         // Handle sparse set removes
         for component_id in self.bundle_info.as_ref().iter_explicit_components() {
             if self.old_archetype.as_ref().contains(component_id) {
@@ -212,14 +225,21 @@ impl<'w> BundleRemover<'w> {
                 if let Some(StorageType::SparseSet) =
                     self.old_archetype.as_ref().get_storage_type(component_id)
                 {
-                    world
+                    let sparse_set = world
                         .storages
                         .sparse_sets
                         .get_mut(component_id)
                         // Set exists because the component existed on the entity
-                        .unwrap()
-                        // If it was already forgotten, it would not be in the set.
-                        .remove(entity);
+                        .unwrap();
+
+                    let maybe_panic =
+                        bevy_utils::catch_unwind_if_available(AssertUnwindSafe(|| {
+                            sparse_set.remove(entity);
+                        }));
+
+                    if panic_payload.is_ok() & maybe_panic.is_err() {
+                        panic_payload = maybe_panic;
+                    }
                 }
             }
         }
@@ -278,6 +298,10 @@ impl<'w> BundleRemover<'w> {
                 }
             };
 
+            if panic_payload.is_ok() & move_result.panic.is_err() {
+                panic_payload = move_result.panic;
+            }
+
             // SAFETY: move_result.new_row is a valid position in new_archetype's table
             let new_location = unsafe {
                 self.new_archetype
@@ -317,7 +341,11 @@ impl<'w> BundleRemover<'w> {
                 .update_existing_location(entity.index(), Some(new_location));
         }
 
-        (new_location, pre_remove_result)
+        BundleRemoveResult {
+            new_location,
+            data: pre_remove_result,
+            panic_payload,
+        }
     }
 }
 

crates/bevy_ecs/src/bundle/spawner.rs

@@ -107,7 +107,8 @@ impl<'w> BundleSpawner<'w> {
             };
             let table_row = table.allocate(entity);
             let location = archetype.allocate(entity, table_row);
-            bundle_info.write_components(
+            // No component existed beforehand, therefore no drop can have panicked
+            let _ = bundle_info.write_components(
                 table,
                 sparse_sets,
                 &SpawnBundleStatus,

crates/bevy_ecs/src/bundle/writer.rs

@@ -126,6 +126,9 @@ impl<'a> BundleWriter<'a> {
     /// Runs with [`RelationshipHookMode::Run`] by default.
     /// Use [`write_with_relationship_hook_insert_mode`](Self::write_with_relationship_hook_insert_mode) if you need more flexibility.
     ///
+    /// # Panics
+    /// Panics if any of the overwritten components panic while being dropped.
+    ///
     /// # Safety
     ///
     /// `entity` must be from the same world that all [`Self::push_component`] or [`Self::push_component_by_id`] calls since the last
@@ -154,17 +157,25 @@ impl<'a> BundleWriter<'a> {
         // - All `component_ids` are from the same world as `entity`
         // - All `component_data_ptrs` are valid types represented by `component_ids`
         unsafe {
+            struct DropGuard<'a>(BundleWriter<'a>);
+            impl Drop for DropGuard<'_> {
+                fn drop(&mut self) {
+                    self.0 .0.component_ids.clear();
+                    self.0 .0.alloc.reset();
+                }
+            }
+            let guard = DropGuard(self);
             entity.insert_by_ids_internal(
-                &self.0.component_ids,
-                self.0
+                &guard.0 .0.component_ids,
+                guard
+                    .0
+                     .0
                     .component_ptrs
                     .drain(..)
                     .map(|ptr| OwningPtr::new(ptr)),
                 relationship_hook_insert_mode,
             );
         }
-        self.0.component_ids.clear();
-        self.0.alloc.reset();
     }
 
     /// Returns true if there are currently no components.