force TLS access to go through mutex

Author: maniwani

crates/bevy_app/src/app.rs

@@ -4,8 +4,8 @@ use bevy_ecs::{
     prelude::*,
     schedule::{IntoSystemConfigs, IntoSystemSetConfigs, ScheduleLabel},
     storage::{
-        ThreadLocalAccessor, ThreadLocalTask, ThreadLocalTaskSendError, ThreadLocalTaskSender,
-        ThreadLocals,
+        ThreadLocalAccessor, ThreadLocalStorage, ThreadLocalTask, ThreadLocalTaskSendError,
+        ThreadLocalTaskSender,
     },
 };
 use bevy_utils::HashMap;
@@ -99,7 +99,7 @@ fn run_once(mut app: App) {
 /// ```
 pub struct App {
     sub_apps: SubApps,
-    tls: ThreadLocals,
+    tls: ThreadLocalStorage,
     /// The function that will manage the app's lifecycle.
     ///
     /// Bevy provides the [`WinitPlugin`] and [`HeadlessPlugin`] for windowed and headless
@@ -155,12 +155,12 @@ impl App {
                 main: SubApp::new(),
                 sub_apps: HashMap::new(),
             },
-            tls: ThreadLocals::new(),
+            tls: ThreadLocalStorage::new(),
             runner: Some(Box::new(run_once)),
         }
     }
 
-    pub fn into_parts(self) -> (SubApps, ThreadLocals, Option<RunnerFn>) {
+    pub fn into_parts(self) -> (SubApps, ThreadLocalStorage, Option<RunnerFn>) {
         let App {
             sub_apps,
             tls,
@@ -170,7 +170,11 @@ impl App {
         (sub_apps, tls, runner)
     }
 
-    pub fn from_parts(sub_apps: SubApps, tls: ThreadLocals, runner: Option<RunnerFn>) -> Self {
+    pub fn from_parts(
+        sub_apps: SubApps,
+        tls: ThreadLocalStorage,
+        runner: Option<RunnerFn>,
+    ) -> Self {
         App {
             sub_apps,
             tls,
@@ -182,16 +186,15 @@ impl App {
     ///
     /// See [`insert_sub_app`](Self::insert_sub_app) for more details.
     pub fn update(&mut self) {
-        let (mut sub_apps, mut tls, runner) = std::mem::take(self).into_parts();
+        let (mut sub_apps, tls, runner) = std::mem::take(self).into_parts();
 
         // create event loop channel
         let (send, recv) = std::sync::mpsc::channel();
 
         // insert TLS accessor
         sub_apps.for_each(|sub_app| {
-            // SAFETY: `tls` is not moved or dropped until `access` has been dropped.
-            let access =
-                unsafe { ThreadLocalAccessor::new(std::ptr::addr_of_mut!(tls), send.clone()) };
+            // SAFETY: `tls` will not be dropped until `access` has been dropped.
+            let access = unsafe { tls.make_accessor(send.clone()) };
             sub_app.world.insert_resource(access);
         });
 
@@ -208,7 +211,7 @@ impl App {
                 let event = recv.recv().unwrap();
                 match event {
                     AppThreadEvent::RunTask(f) => {
-                        f(&mut tls);
+                        f(&mut *tls.lock());
                     }
                     AppThreadEvent::Exit(sub_apps) => {
                         handle.join();
@@ -475,7 +478,7 @@ impl App {
     ///     .insert_non_send_resource(MyCounter { counter: 0 });
     /// ```
     pub fn insert_non_send_resource<R: ThreadLocalResource>(&mut self, resource: R) -> &mut Self {
-        self.tls.insert_resource(resource);
+        self.tls.lock().insert_resource(resource);
         self
     }
 
@@ -485,7 +488,7 @@ impl App {
     /// If the [`Default`] trait is implemented, the [`FromWorld`] trait will use
     /// the [`Default::default`] method to initialize the [`Resource`].
     pub fn init_non_send_resource<R: 'static + FromWorld>(&mut self) -> &mut Self {
-        self.tls.init_resource::<R>();
+        self.tls.lock().init_resource::<R>();
         self
     }
 

crates/bevy_app/src/schedule_runner.rs

@@ -102,17 +102,15 @@ impl Plugin for ScheduleRunnerPlugin {
                         Ok(Duration::ZERO)
                     };
 
-                    let (mut sub_apps, mut tls, _) = app.into_parts();
+                    let (mut sub_apps, tls, _) = app.into_parts();
 
                     // create event loop channel
                     let (send, recv) = std::sync::mpsc::channel();
 
                     // insert TLS accessor
                     sub_apps.for_each(|sub_app| {
-                        // SAFETY: `tls` is not moved or dropped until `access` has been dropped.
-                        let access = unsafe {
-                            ThreadLocalAccessor::new(std::ptr::addr_of_mut!(tls), send.clone())
-                        };
+                        // SAFETY: `tls` will not be dropped until `access` has been dropped.
+                        let access = unsafe { tls.make_accessor(send.clone()) };
                         sub_app.world.insert_resource(access);
                     });
 
@@ -133,7 +131,7 @@ impl Plugin for ScheduleRunnerPlugin {
                             let event = recv.recv().unwrap();
                             match event {
                                 AppThreadEvent::RunTask(f) => {
-                                    f(&mut tls);
+                                    f(&mut *tls.lock());
                                 }
                                 AppThreadEvent::Exit(sub_apps) => {
                                     handle.join();