bevy_app changes

maniwani · maniwani · commit b4420b3a0315 · 2023-09-21T14:47:35.000-07:00
diff --git a/crates/bevy_app/src/app.rs b/crates/bevy_app/src/app.rs
@@ -1,10 +1,15 @@
-use crate::{Main, MainSchedulePlugin, PlaceholderPlugin, Plugin, Plugins, SubApp, SubApps};
+use crate::{
+    app_thread_channel, AppEvent, Main, MainSchedulePlugin, PlaceholderPlugin, Plugin, Plugins,
+    SubApp, SubApps,
+};
 pub use bevy_derive::AppLabel;
 use bevy_ecs::schedule::ScheduleBuildSettings;
 use bevy_ecs::{
     prelude::*,
     schedule::{IntoSystemConfigs, IntoSystemSetConfigs, ScheduleLabel},
+    storage::ThreadLocalStorage,
 };
+
 #[cfg(feature = "trace")]
 use bevy_utils::tracing::info_span;
 use bevy_utils::{tracing::debug, HashMap};
@@ -50,7 +55,10 @@ pub(crate) enum AppError {
 /// }
 /// ```
 pub struct App {
-    pub(crate) sub_apps: SubApps,
+    #[doc(hidden)]
+    pub sub_apps: SubApps,
+    #[doc(hidden)]
+    pub tls: ThreadLocalStorage,
     /// The function that will manage the app's lifecycle.
     ///
     /// Bevy provides the [`WinitPlugin`] and [`ScheduleRunnerPlugin`] for windowed and headless
@@ -106,17 +114,89 @@ impl App {
                 main: SubApp::new(),
                 sub_apps: HashMap::new(),
             },
+            tls: ThreadLocalStorage::new(),
             runner: Some(Box::new(run_once)),
         }
     }
 
+    /// Disassembles the [`App`] and returns its individual parts.
+    pub fn into_parts(self) -> (SubApps, ThreadLocalStorage, Option<RunnerFn>) {
+        let Self {
+            sub_apps,
+            tls,
+            runner,
+        } = self;
+
+        (sub_apps, tls, runner)
+    }
+
+    /// Returns an [`App`] assembled from the given individual parts.
+    pub fn from_parts(
+        sub_apps: SubApps,
+        tls: ThreadLocalStorage,
+        runner: Option<RunnerFn>,
+    ) -> Self {
+        App {
+            sub_apps,
+            tls,
+            runner,
+        }
+    }
+
     /// Runs the default schedules of all sub-apps (starting with the "main" app) once.
     pub fn update(&mut self) {
         if self.is_building_plugins() {
             panic!("App::update() was called while a plugin was building.");
         }
 
-        self.sub_apps.update();
+        // disassemble
+        let (mut sub_apps, tls, runner) = std::mem::take(self).into_parts();
+
+        // create channel
+        let (send, recv) = app_thread_channel();
+
+        // insert channel
+        sub_apps
+            .iter_mut()
+            .for_each(|sub_app| tls.insert_channel(sub_app.world_mut(), send.clone()));
+
+        #[cfg(not(target_arch = "wasm32"))]
+        {
+            // Move sub-apps to another thread and run an event loop in this thread.
+            let thread = std::thread::spawn(move || {
+                sub_apps.update();
+                send.send(AppEvent::Exit(Box::new(sub_apps))).unwrap();
+            });
+
+            loop {
+                let event = recv.recv().unwrap();
+                match event {
+                    AppEvent::Task(f) => {
+                        f(&mut tls.lock());
+                    }
+                    AppEvent::Exit(boxed) => {
+                        // SAFETY: `Box::<T>::into_raw` returns a pointer that is properly aligned
+                        // and points to an initialized value of `T`.
+                        sub_apps = unsafe { std::ptr::read(Box::into_raw(boxed)) };
+                        thread.join().unwrap();
+                        break;
+                    }
+                }
+            }
+        }
+
+        #[cfg(target_arch = "wasm32")]
+        {
+            sub_apps.update();
+        }
+
+        // remove channel
+        sub_apps
+            .iter_mut()
+            .for_each(|sub_app| tls.remove_channel(sub_app.world_mut()));
+
+        // reassemble
+        *self = App::from_parts(sub_apps, tls, runner);
     }
 
     /// Runs the [`App`] by calling its [runner](Self::set_runner).
@@ -387,30 +467,31 @@ impl App {
     /// of the same type.
     ///
     /// There is also an [`init_non_send_resource`](Self::init_non_send_resource) for
-    /// resources that implement [`Default`]
+    /// resources that implement [`Default`].
     ///
     /// # Examples
     ///
     /// ```
     /// # use bevy_app::prelude::*;
     /// # use bevy_ecs::prelude::*;
     /// #
+    /// #[derive(ThreadLocalResource)]
     /// struct MyCounter {
     ///     counter: usize,
     /// }
     ///
     /// App::new()
     ///     .insert_non_send_resource(MyCounter { counter: 0 });
     /// ```
-    pub fn insert_non_send_resource<R: 'static>(&mut self, resource: R) -> &mut Self {
-        self.world_mut().insert_non_send_resource(resource);
+    pub fn insert_non_send_resource<R: ThreadLocalResource>(&mut self, resource: R) -> &mut Self {
+        self.tls.lock().insert_resource(resource);
         self
     }
 
     /// Inserts the [`!Send`](Send) resource into the app, initialized with its default value,
     /// if there is no existing instance of `R`.
-    pub fn init_non_send_resource<R: 'static + Default>(&mut self) -> &mut Self {
-        self.world_mut().init_non_send_resource::<R>();
+    pub fn init_non_send_resource<R: ThreadLocalResource + Default>(&mut self) -> &mut Self {
+        self.tls.lock().init_resource::<R>();
         self
     }
 
@@ -692,13 +773,55 @@ impl App {
 type RunnerFn = Box<dyn FnOnce(App)>;
 
 fn run_once(mut app: App) {
+    if app.is_building_plugins() {
+        panic!("App::update() was called while a plugin was building.");
+    }
+
+    // TODO: rework app setup
+    // create channel
+    let (send, recv) = app_thread_channel();
+    // insert channel
+    app.sub_apps
+        .iter_mut()
+        .for_each(|sub_app| app.tls.insert_channel(sub_app.world_mut(), send.clone()));
+
+    // wait for plugins to finish setting up
     while !app.is_ready() {
         #[cfg(not(target_arch = "wasm32"))]
         bevy_tasks::tick_global_task_pools_on_main_thread();
     }
     app.finish();
     app.cleanup();
-    app.update();
+
+    // disassemble
+    let (mut sub_apps, tls, _) = app.into_parts();
+
+    #[cfg(not(target_arch = "wasm32"))]
+    {
+        // Move sub-apps to another thread and run an event loop in this thread.
+        let thread = std::thread::spawn(move || {
+            sub_apps.update();
+            send.send(AppEvent::Exit(Box::new(sub_apps))).unwrap();
+        });
+
+        loop {
+            let event = recv.recv().unwrap();
+            match event {
+                AppEvent::Task(f) => {
+                    f(&mut tls.lock());
+                }
+                AppEvent::Exit(_) => {
+                    thread.join().unwrap();
+                    break;
+                }
+            }
+        }
+    }
+
+    #[cfg(target_arch = "wasm32")]
+    {
+        sub_apps.update();
+    }
 }
 
 /// An event that indicates the [`App`]  should exit. If one or more of these are present at the
diff --git a/crates/bevy_app/src/events.rs b/crates/bevy_app/src/events.rs
@@ -0,0 +1,64 @@
+use bevy_ecs::storage::{ThreadLocalTask, ThreadLocalTaskSendError, ThreadLocalTaskSender};
+use std::sync::mpsc::{channel, Receiver, Sender};
+
+/// Events an [`App`](crate::App) can send to another thread when using a multi-threaded runner.
+pub enum AppEvent {
+    /// The app has sent a task with access to [`ThreadLocals`](bevy_ecs::prelude::ThreadLocals).
+    Task(ThreadLocalTask),
+    /// The app has exited.
+    Exit(Box<crate::SubApps>),
+}
+
+/// The sender half of an [`app_thread_channel`].
+#[derive(Clone)]
+pub struct AppEventSender(Sender<AppEvent>);
+
+/// Constructs a new asynchronous channel for passing [`AppEvent`] instances
+/// to an event loop and returns the sender and receiver halves.
+pub fn app_thread_channel() -> (AppEventSender, AppEventReceiver) {
+    let (send, recv) = channel();
+    (AppEventSender(send), AppEventReceiver(recv))
+}
+
+impl std::ops::Deref for AppEventSender {
+    type Target = Sender<AppEvent>;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl std::ops::DerefMut for AppEventSender {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl ThreadLocalTaskSender for AppEventSender {
+    fn send_task(
+        &mut self,
+        task: ThreadLocalTask,
+    ) -> Result<(), ThreadLocalTaskSendError<ThreadLocalTask>> {
+        self.send(AppEvent::Task(task)).map_err(|error| {
+            let AppEvent::Task(task) = error.0 else {
+                unreachable!()
+            };
+            ThreadLocalTaskSendError(task)
+        })
+    }
+}
+
+/// The receiver-half of an [`app_thread_channel`].
+pub struct AppEventReceiver(Receiver<AppEvent>);
+
+impl std::ops::Deref for AppEventReceiver {
+    type Target = Receiver<AppEvent>;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl std::ops::DerefMut for AppEventReceiver {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
diff --git a/crates/bevy_app/src/lib.rs b/crates/bevy_app/src/lib.rs
@@ -4,6 +4,7 @@
 #![allow(clippy::type_complexity)]
 
 mod app;
+mod events;
 mod main_schedule;
 mod plugin;
 mod plugin_group;
@@ -15,6 +16,7 @@ pub mod ci_testing;
 
 pub use app::*;
 pub use bevy_derive::DynamicPlugin;
+pub use events::*;
 pub use main_schedule::*;
 pub use plugin::*;
 pub use plugin_group::*;
diff --git a/crates/bevy_app/src/schedule_runner.rs b/crates/bevy_app/src/schedule_runner.rs
