Fix Plugin::build detection (#8103)

crates/bevy_app/src/app.rs

@@ -12,7 +12,10 @@ use bevy_ecs::{
     },
 };
 use bevy_utils::{tracing::debug, HashMap, HashSet};
-use std::fmt::Debug;
+use std::{
+    fmt::Debug,
+    panic::{catch_unwind, resume_unwind, AssertUnwindSafe},
+};
 
 #[cfg(feature = "trace")]
 use bevy_utils::tracing::info_span;
@@ -80,8 +83,8 @@ pub struct App {
     sub_apps: HashMap<AppLabelId, SubApp>,
     plugin_registry: Vec<Box<dyn Plugin>>,
     plugin_name_added: HashSet<String>,
-    /// A private marker to prevent incorrect calls to `App::run()` from `Plugin::build()`
-    is_building_plugin: bool,
+    /// A private counter to prevent incorrect calls to `App::run()` from `Plugin::build()`
+    building_plugin_depth: usize,
 }
 
 impl Debug for App {
@@ -228,7 +231,7 @@ impl App {
             plugin_name_added: Default::default(),
             default_schedule_label: Box::new(CoreSchedule::Main),
             outer_schedule_label: Box::new(CoreSchedule::Outer),
-            is_building_plugin: false,
+            building_plugin_depth: 0,
         }
     }
 
@@ -291,8 +294,8 @@ impl App {
         let _bevy_app_run_span = info_span!("bevy_app").entered();
 
         let mut app = std::mem::replace(self, App::empty());
-        if app.is_building_plugin {
-            panic!("App::run() was called from within Plugin::Build(), which is not allowed.");
+        if app.building_plugin_depth > 0 {
+            panic!("App::run() was called from within Plugin::build(), which is not allowed.");
         }
 
         Self::setup(&mut app);
@@ -765,9 +768,12 @@ impl App {
                 plugin_name: plugin.name().to_string(),
             })?;
         }
-        self.is_building_plugin = true;
-        plugin.build(self);
-        self.is_building_plugin = false;
+        self.building_plugin_depth += 1;
+        let result = catch_unwind(AssertUnwindSafe(|| plugin.build(self)));
+        self.building_plugin_depth -= 1;
+        if let Err(payload) = result {
+            resume_unwind(payload);
+        }
         self.plugin_registry.push(plugin);
         Ok(self)
     }
@@ -1071,9 +1077,13 @@ mod tests {
     #[should_panic]
     fn cant_call_app_run_from_plugin_build() {
         struct PluginRun;
+        struct InnerPlugin;
+        impl Plugin for InnerPlugin {
+            fn build(&self, _: &mut crate::App) {}
+        }
         impl Plugin for PluginRun {
             fn build(&self, app: &mut crate::App) {
-                app.run();
+                app.add_plugin(InnerPlugin).run();
             }
         }
         App::new().add_plugin(PluginRun);