[Merged by Bors] - Stageless: add a method to scope to always run a task on the scope thread

crates/bevy_ecs/src/schedule/executor_parallel.rs

@@ -1,15 +1,12 @@
-use std::sync::Arc;
-
-use crate as bevy_ecs;
 use crate::{
     archetype::ArchetypeComponentId,
     query::Access,
     schedule::{ParallelSystemExecutor, SystemContainer},
-    system::Resource,
+    schedule_v3::MainThreadExecutor,
     world::World,
 };
 use async_channel::{Receiver, Sender};
-use bevy_tasks::{ComputeTaskPool, Scope, TaskPool, ThreadExecutor};
+use bevy_tasks::{ComputeTaskPool, Scope, TaskPool};
 #[cfg(feature = "trace")]
 use bevy_utils::tracing::Instrument;
 use event_listener::Event;
@@ -18,16 +15,6 @@ use fixedbitset::FixedBitSet;
 #[cfg(test)]
 use scheduling_event::*;
 
-/// New-typed [`ThreadExecutor`] [`Resource`] that is used to run systems on the main thread
-#[derive(Resource, Default, Clone)]
-pub struct MainThreadExecutor(pub Arc<ThreadExecutor<'static>>);
-
-impl MainThreadExecutor {
-    pub fn new() -> Self {
-        MainThreadExecutor(Arc::new(ThreadExecutor::new()))
-    }
-}
-
 struct SystemSchedulingMetadata {
     /// Used to signal the system's task to start the system.
     start: Event,
@@ -138,7 +125,10 @@ impl ParallelSystemExecutor for ParallelExecutor {
             }
         }
 
-        let thread_executor = world.get_resource::<MainThreadExecutor>().map(|e| &*e.0);
+        let thread_executor = world
+            .get_resource::<MainThreadExecutor>()
+            .map(|e| e.0.clone());
+        let thread_executor = thread_executor.as_deref();
 
         ComputeTaskPool::init(TaskPool::default).scope_with_executor(
             false,
@@ -256,7 +246,7 @@ impl ParallelExecutor {
                 if system_data.is_send {
                     scope.spawn(task);
                 } else {
-                    scope.spawn_on_scope(task);
+                    scope.spawn_on_external(task);
                 }
 
                 #[cfg(test)]
@@ -291,7 +281,7 @@ impl ParallelExecutor {
                 if system_data.is_send {
                     scope.spawn(task);
                 } else {
-                    scope.spawn_on_scope(task);
+                    scope.spawn_on_external(task);
                 }
             }
         }

crates/bevy_ecs/src/schedule_v3/executor/mod.rs

@@ -2,7 +2,7 @@ mod multi_threaded;
 mod simple;
 mod single_threaded;
 
-pub use self::multi_threaded::MultiThreadedExecutor;
+pub use self::multi_threaded::{MainThreadExecutor, MultiThreadedExecutor};
 pub use self::simple::SimpleExecutor;
 pub use self::single_threaded::SingleThreadedExecutor;
 

crates/bevy_ecs/src/schedule_v3/executor/multi_threaded.rs

@@ -1,14 +1,16 @@
-use bevy_tasks::{ComputeTaskPool, Scope, TaskPool};
+use bevy_tasks::{ComputeTaskPool, Scope, TaskPool, ThreadExecutor};
 use bevy_utils::default;
 use bevy_utils::syncunsafecell::SyncUnsafeCell;
 #[cfg(feature = "trace")]
 use bevy_utils::tracing::{info_span, Instrument};
+use std::sync::Arc;
 
 use async_channel::{Receiver, Sender};
 use fixedbitset::FixedBitSet;
 
 use crate::{
     archetype::ArchetypeComponentId,
+    prelude::Resource,
     query::Access,
     schedule_v3::{
         is_apply_system_buffers, BoxedCondition, ExecutorKind, SystemExecutor, SystemSchedule,
@@ -17,6 +19,8 @@ use crate::{
     world::World,
 };
 
+use crate as bevy_ecs;
+
 /// A funky borrow split of [`SystemSchedule`] required by the [`MultiThreadedExecutor`].
 struct SyncUnsafeSchedule<'a> {
     systems: &'a [SyncUnsafeCell<BoxedSystem>],
@@ -145,47 +149,56 @@ impl SystemExecutor for MultiThreadedExecutor {
             }
         }
 
+        let thread_executor = world
+            .get_resource::<MainThreadExecutor>()
+            .map(|e| e.0.clone());
+        let thread_executor = thread_executor.as_deref();
+
         let world = SyncUnsafeCell::from_mut(world);
         let SyncUnsafeSchedule {
             systems,
             mut conditions,
         } = SyncUnsafeSchedule::new(schedule);
 
-        ComputeTaskPool::init(TaskPool::default).scope(|scope| {
-            // the executor itself is a `Send` future so that it can run
-            // alongside systems that claim the local thread
-            let executor = async {
-                while self.num_completed_systems < num_systems {
-                    // SAFETY: self.ready_systems does not contain running systems
-                    unsafe {
-                        self.spawn_system_tasks(scope, systems, &mut conditions, world);
-                    }
-
-                    if self.num_running_systems > 0 {
-                        // wait for systems to complete
-                        let index = self
-                            .receiver
-                            .recv()
-                            .await
-                            .unwrap_or_else(|error| unreachable!("{}", error));
+        ComputeTaskPool::init(TaskPool::default).scope_with_executor(
+            false,
+            thread_executor,
+            |scope| {
+                // the executor itself is a `Send` future so that it can run
+                // alongside systems that claim the local thread
+                let executor = async {
+                    while self.num_completed_systems < num_systems {
+                        // SAFETY: self.ready_systems does not contain running systems
+                        unsafe {
+                            self.spawn_system_tasks(scope, systems, &mut conditions, world);
+                        }
 
-                        self.finish_system_and_signal_dependents(index);
+                        if self.num_running_systems > 0 {
+                            // wait for systems to complete
+                            let index = self
+                                .receiver
+                                .recv()
+                                .await
+                                .unwrap_or_else(|error| unreachable!("{}", error));
 
-                        while let Ok(index) = self.receiver.try_recv() {
                             self.finish_system_and_signal_dependents(index);
-                        }
 
-                        self.rebuild_active_access();
+                            while let Ok(index) = self.receiver.try_recv() {
+                                self.finish_system_and_signal_dependents(index);
+                            }
+
+                            self.rebuild_active_access();
+                        }
                     }
-                }
-            };
+                };
 
-            #[cfg(feature = "trace")]
-            let executor_span = info_span!("schedule_task");
-            #[cfg(feature = "trace")]
-            let executor = executor.instrument(executor_span);
-            scope.spawn(executor);
-        });
+                #[cfg(feature = "trace")]
+                let executor_span = info_span!("schedule_task");
+                #[cfg(feature = "trace")]
+                let executor = executor.instrument(executor_span);
+                scope.spawn(executor);
+            },
+        );
 
         // Do one final apply buffers after all systems have completed
         // SAFETY: all systems have completed, and so no outstanding accesses remain
@@ -437,7 +450,7 @@ impl MultiThreadedExecutor {
             scope.spawn(task);
         } else {
             self.local_thread_running = true;
-            scope.spawn_on_scope(task);
+            scope.spawn_on_external(task);
         }
     }
 
@@ -574,3 +587,13 @@ fn evaluate_and_fold_conditions(conditions: &mut [BoxedCondition], world: &World
         })
         .fold(true, |acc, res| acc && res)
 }
+
+/// New-typed [`ThreadExecutor`] [`Resource`] that is used to run systems on the main thread
+#[derive(Resource, Default, Clone)]
+pub struct MainThreadExecutor(pub Arc<ThreadExecutor<'static>>);
+
+impl MainThreadExecutor {
+    pub fn new() -> Self {
+        MainThreadExecutor(Arc::new(ThreadExecutor::new()))
+    }
+}

crates/bevy_render/src/pipelined_rendering.rs

@@ -2,7 +2,8 @@ use async_channel::{Receiver, Sender};
 
 use bevy_app::{App, AppLabel, Plugin, SubApp};
 use bevy_ecs::{
-    schedule::{MainThreadExecutor, StageLabel, SystemStage},
+    schedule::{StageLabel, SystemStage},
+    schedule_v3::MainThreadExecutor,
     system::Resource,
     world::{Mut, World},
 };

crates/bevy_tasks/src/single_threaded_task_pool.rs

@@ -178,17 +178,28 @@ pub struct Scope<'scope, 'env: 'scope, T> {
 }
 
 impl<'scope, 'env, T: Send + 'env> Scope<'scope, 'env, T> {
-    /// Spawns a scoped future onto the thread-local executor. The scope *must* outlive
+    /// Spawns a scoped future onto the executor. The scope *must* outlive
     /// the provided future. The results of the future will be returned as a part of
     /// [`TaskPool::scope`]'s return value.
     ///
-    /// On the single threaded task pool, it just calls [`Scope::spawn_local`].
+    /// On the single threaded task pool, it just calls [`Scope::spawn_on_scope`].
     ///
     /// For more information, see [`TaskPool::scope`].
     pub fn spawn<Fut: Future<Output = T> + 'env>(&self, f: Fut) {
         self.spawn_on_scope(f);
     }
 
+    /// Spawns a scoped future onto the executor. The scope *must* outlive
+    /// the provided future. The results of the future will be returned as a part of
+    /// [`TaskPool::scope`]'s return value.
+    ///
+    /// On the single threaded task pool, it just calls [`Scope::spawn_on_scope`].
+    ///
+    /// For more information, see [`TaskPool::scope`].
+    pub fn spawn_on_external<Fut: Future<Output = T> + 'env>(&self, f: Fut) {
+        self.spawn_on_scope(f);
+    }
+
     /// Spawns a scoped future that runs on the thread the scope called from. The
     /// scope *must* outlive the provided future. The results of the future will be
     /// returned as a part of [`TaskPool::scope`]'s return value.