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,
-    ) {
+    ) -> Option<Box<dyn Any + Send>> {
+        let mut panic = None;
+
         // 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)).err();
+            if panic.is_none() {
+                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::{
@@ -137,6 +138,7 @@ impl<'w> BundleInserter<'w> {
         EntityLocation,
         &'a mut SparseSets,
         &'a mut Table,
+        Option<Box<dyn Any + Send>>,
     ) {
         // SAFETY: All components in the bundle are guaranteed to exist in the World
         // as they must be initialized before creating the BundleInfo.
@@ -191,7 +193,7 @@ impl<'w> BundleInserter<'w> {
                     )
                 };
 
-                (&*archetype, location, sparse_sets, table)
+                (&*archetype, location, sparse_sets, table, None)
             }
             ArchetypeMoveType::NewArchetypeSameTable { new_archetype } => {
                 let new_archetype = new_archetype.as_mut();
@@ -224,7 +226,7 @@ impl<'w> BundleInserter<'w> {
                 let new_location = new_archetype.allocate(entity, result.table_row);
                 entities.update_existing_location(entity.index(), Some(new_location));
 
-                (&*new_archetype, new_location, sparse_sets, table)
+                (&*new_archetype, new_location, sparse_sets, table, None)
             }
             ArchetypeMoveType::NewArchetypeNewTable { new_archetype } => {
                 let new_archetype = new_archetype.as_mut();
@@ -310,11 +312,14 @@ impl<'w> BundleInserter<'w> {
                     new_location,
                     sparse_sets,
                     move_result.new_table,
+                    move_result.panic,
                 )
             }
         }
     }
 
+    /// 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,24 +339,27 @@ impl<'w> BundleInserter<'w> {
         insert_mode: InsertMode,
         caller: MaybeLocation,
         relationship_hook_mode: RelationshipHookMode,
-    ) -> EntityLocation {
+    ) -> (EntityLocation, Option<Box<dyn Any + Send>>) {
         let archetype_after_insert = self.archetype_after_insert.as_ref();
 
+        let panic_payload;
+
         let (new_archetype, new_location) = {
             // Non-generic prelude extracted to improve compile time by minimizing monomorphized code.
-            let (new_archetype, new_location, sparse_sets, table) = Self::before_insert(
-                entity,
-                location,
-                insert_mode,
-                caller,
-                relationship_hook_mode,
-                self.archetype,
-                archetype_after_insert,
-                &self.world,
-                &mut self.archetype_move_type,
-            );
+            let (new_archetype, new_location, sparse_sets, table, archetype_move_panic_payload) =
+                Self::before_insert(
+                    entity,
+                    location,
+                    insert_mode,
+                    caller,
+                    relationship_hook_mode,
+                    self.archetype,
+                    archetype_after_insert,
+                    &self.world,
+                    &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,6 +372,8 @@ impl<'w> BundleInserter<'w> {
                 caller,
             );
 
+            panic_payload = archetype_move_panic_payload.or(maybe_panic);
+
             (new_archetype, new_location)
         };
 
@@ -382,7 +392,7 @@ impl<'w> BundleInserter<'w> {
             deferred_world,
         );
 
-        new_location
+        (new_location, panic_payload)
     }
 
     // A non-generic postlude to insert used to minimize duplicated monomorphized code.

crates/bevy_ecs/src/bundle/remove.rs

@@ -1,5 +1,6 @@
 use alloc::vec::Vec;
 use bevy_ptr::ConstNonNull;
+use core::panic::AssertUnwindSafe;
 use core::ptr::NonNull;
 
 use crate::{
@@ -26,6 +27,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: Option<alloc::boxed::Box<dyn core::any::Any + Send>>,
+}
+
 impl<'w> BundleRemover<'w> {
     /// Creates a new [`BundleRemover`], if such a remover would do anything.
     ///
@@ -133,7 +140,7 @@ impl<'w> BundleRemover<'w> {
             &Components,
             &[ComponentId],
         ) -> (bool, T),
-    ) -> (EntityLocation, T) {
+    ) -> BundleRemoveResult<T> {
         // Hooks
         // SAFETY: all bundle components exist in World
         unsafe {
@@ -202,6 +209,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 = None;
+
         // 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 +224,20 @@ 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)
+                        }))
+                        .err();
+
+                    panic_payload = panic_payload.or(maybe_panic);
                 }
             }
         }
@@ -278,6 +296,12 @@ impl<'w> BundleRemover<'w> {
                 }
             };
 
+            if let Some(payload) = move_result.panic
+                && panic_payload.is_none()
+            {
+                panic_payload = Some(payload);
+            }
+
             // 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/writer.rs

@@ -6,7 +6,7 @@ use crate::{
 use alloc::vec::Vec;
 use bevy_ptr::OwningPtr;
 use bumpalo::Bump;
-use core::{alloc::Layout, ptr::NonNull};
+use core::{alloc::Layout, panic::AssertUnwindSafe, ptr::NonNull};
 
 /// Enables pushing components to internal scratch space (uses a bump allocator), which can then be
 /// written as a dynamic bundle. The contents are cleared after each write and the allocated scratch
@@ -123,6 +123,9 @@ impl<'a> BundleWriter<'a> {
 
     /// Writes the current contents of the bundle to the given `entity` and clears the scratch space.
     ///
+    /// # 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`] calls since the last
@@ -132,17 +135,20 @@ 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 {
-            entity.insert_by_ids_internal(
-                &self.0.component_ids,
-                self.0
-                    .component_ptrs
-                    .drain(..)
-                    .map(|ptr| OwningPtr::new(ptr)),
-                RelationshipHookMode::Run,
-            );
+            let maybe_panic = bevy_utils::catch_unwind_if_available(AssertUnwindSafe(|| {
+                entity.insert_by_ids_internal(
+                    &self.0.component_ids,
+                    self.0
+                        .component_ptrs
+                        .drain(..)
+                        .map(|ptr| OwningPtr::new(ptr)),
+                    RelationshipHookMode::Run,
+                );
+            }));
+            self.0.component_ids.clear();
+            self.0.alloc.reset();
+            bevy_utils::resume_caught_unwind(maybe_panic.err());
         }
-        self.0.component_ids.clear();
-        self.0.alloc.reset();
     }
 
     /// Returns true if there are currently no components.

crates/bevy_ecs/src/schedule/executor/single_threaded.rs

@@ -131,7 +131,7 @@ impl SystemExecutor for SingleThreadedExecutor {
                 continue;
             }
 
-            let f = AssertUnwindSafe(|| {
+            let maybe_panic = bevy_utils::catch_unwind_if_available(AssertUnwindSafe(|| {
                 if let Err(RunSystemError::Failed(err)) =
                     __rust_begin_short_backtrace::run_without_applying_deferred(system, world)
                 {
@@ -143,21 +143,15 @@ impl SystemExecutor for SingleThreadedExecutor {
                         },
                     );
                 }
-            });
+            }))
+            .err();
 
             #[cfg(feature = "std")]
             #[expect(clippy::print_stderr, reason = "Allowed behind `std` feature gate.")]
-            {
-                if let Err(payload) = std::panic::catch_unwind(f) {
-                    eprintln!("Encountered a panic in system `{}`!", system.name());
-                    std::panic::resume_unwind(payload);
-                }
-            }
-
-            #[cfg(not(feature = "std"))]
-            {
-                (f)();
+            if maybe_panic.is_some() {
+                eprintln!("Encountered a panic in system `{}`!", system.name());
             }
+            bevy_utils::resume_caught_unwind(maybe_panic);
 
             self.unapplied_systems.insert(system_index);
         }