diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index e9f56a58a8ec9..739a1d9d30f31 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -1191,3 +1191,136 @@ impl Default for MainThreadValidator { } } } + +#[cfg(test)] +mod tests { + use super::World; + use bevy_ecs_macros::Component; + use std::{ + panic, + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, + }, + }; + + // For bevy_ecs_macros + use crate as bevy_ecs; + + type ID = u8; + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum DropLogItem { + Create(ID), + Drop(ID), + } + + #[derive(Component)] + struct MayPanicInDrop { + drop_log: Arc>>, + expected_panic_flag: Arc, + should_panic: bool, + id: u8, + } + + impl MayPanicInDrop { + fn new( + drop_log: &Arc>>, + expected_panic_flag: &Arc, + should_panic: bool, + id: u8, + ) -> Self { + println!("creating component with id {}", id); + drop_log.lock().unwrap().push(DropLogItem::Create(id)); + + Self { + drop_log: Arc::clone(drop_log), + expected_panic_flag: Arc::clone(expected_panic_flag), + should_panic, + id, + } + } + } + + impl Drop for MayPanicInDrop { + fn drop(&mut self) { + println!("dropping component with id {}", self.id); + + { + let mut drop_log = self.drop_log.lock().unwrap(); + drop_log.push(DropLogItem::Drop(self.id)); + // Don't keep the mutex while panicking, or we'll poison it. + drop(drop_log); + } + + if self.should_panic { + self.expected_panic_flag.store(true, Ordering::SeqCst); + panic!("testing what happens on panic inside drop"); + } + } + } + + struct DropTestHelper { + drop_log: Arc>>, + /// Set to `true` right before we intentionally panic, so that if we get + /// a panic, we know if it was intended or not. + expected_panic_flag: Arc, + } + + impl DropTestHelper { + pub fn new() -> Self { + Self { + drop_log: Arc::new(Mutex::new(Vec::::new())), + expected_panic_flag: Arc::new(AtomicBool::new(false)), + } + } + + pub fn make_component(&self, should_panic: bool, id: ID) -> MayPanicInDrop { + MayPanicInDrop::new(&self.drop_log, &self.expected_panic_flag, should_panic, id) + } + + pub fn finish(self, panic_res: std::thread::Result<()>) -> Vec { + let drop_log = Arc::try_unwrap(self.drop_log) + .unwrap() + .into_inner() + .unwrap(); + let expected_panic_flag = self.expected_panic_flag.load(Ordering::SeqCst); + + if !expected_panic_flag { + match panic_res { + Ok(()) => panic!("Expected a panic but it didn't happen"), + Err(e) => panic::resume_unwind(e), + } + } + + drop_log + } + } + + #[test] + fn panic_while_overwriting_component() { + let helper = DropTestHelper::new(); + + let res = panic::catch_unwind(|| { + let mut world = World::new(); + world + .spawn() + .insert(helper.make_component(true, 0)) + .insert(helper.make_component(false, 1)); + + println!("Done inserting! Dropping world..."); + }); + + let drop_log = helper.finish(res); + + assert_eq!( + &*drop_log, + [ + DropLogItem::Create(0), + DropLogItem::Create(1), + DropLogItem::Drop(0), + DropLogItem::Drop(1) + ] + ); + } +}