Global TaskPool API improvements (#10008)

# Objective

Reduce code duplication and improve APIs of Bevy's [global
taskpools](https://github.com/bevyengine/bevy/blob/main/crates/bevy_tasks/src/usages.rs).

## Solution

- As all three of the global taskpools have identical implementations
and only differ in their identifiers, this PR moves the implementation
into a macro to reduce code duplication.
- The `init` method is renamed to `get_or_init` to more accurately
reflect what it really does.
- Add a new `try_get` method that just returns `None` when the pool is
uninitialized, to complement the other getter methods.
- Minor documentation improvements to accompany the above changes.

---

## Changelog

- Added a new `try_get` method to the global TaskPools
- The global TaskPools' `init` method has been renamed to `get_or_init`
for clarity
- Documentation improvements

## Migration Guide

- Uses of `ComputeTaskPool::init`, `AsyncComputeTaskPool::init` and
`IoTaskPool::init` should be changed to `::get_or_init`.

benches/benches/bevy_ecs/iteration/heavy_compute.rs

@@ -20,7 +20,7 @@ pub fn heavy_compute(c: &mut Criterion) {
     group.warm_up_time(std::time::Duration::from_millis(500));
     group.measurement_time(std::time::Duration::from_secs(4));
     group.bench_function("base", |b| {
-        ComputeTaskPool::init(TaskPool::default);
+        ComputeTaskPool::get_or_init(TaskPool::default);
 
         let mut world = World::default();
 

crates/bevy_core/src/task_pool_options.rs

@@ -107,7 +107,7 @@ impl TaskPoolOptions {
             trace!("IO Threads: {}", io_threads);
             remaining_threads = remaining_threads.saturating_sub(io_threads);
 
-            IoTaskPool::init(|| {
+            IoTaskPool::get_or_init(|| {
                 TaskPoolBuilder::default()
                     .num_threads(io_threads)
                     .thread_name("IO Task Pool".to_string())
@@ -124,7 +124,7 @@ impl TaskPoolOptions {
             trace!("Async Compute Threads: {}", async_compute_threads);
             remaining_threads = remaining_threads.saturating_sub(async_compute_threads);
 
-            AsyncComputeTaskPool::init(|| {
+            AsyncComputeTaskPool::get_or_init(|| {
                 TaskPoolBuilder::default()
                     .num_threads(async_compute_threads)
                     .thread_name("Async Compute Task Pool".to_string())
@@ -141,7 +141,7 @@ impl TaskPoolOptions {
 
             trace!("Compute Threads: {}", compute_threads);
 
-            ComputeTaskPool::init(|| {
+            ComputeTaskPool::get_or_init(|| {
                 TaskPoolBuilder::default()
                     .num_threads(compute_threads)
                     .thread_name("Compute Task Pool".to_string())

crates/bevy_ecs/src/lib.rs

@@ -400,7 +400,7 @@ mod tests {
 
     #[test]
     fn par_for_each_dense() {
-        ComputeTaskPool::init(TaskPool::default);
+        ComputeTaskPool::get_or_init(TaskPool::default);
         let mut world = World::new();
         let e1 = world.spawn(A(1)).id();
         let e2 = world.spawn(A(2)).id();
@@ -423,7 +423,7 @@ mod tests {
 
     #[test]
     fn par_for_each_sparse() {
-        ComputeTaskPool::init(TaskPool::default);
+        ComputeTaskPool::get_or_init(TaskPool::default);
         let mut world = World::new();
         let e1 = world.spawn(SparseStored(1)).id();
         let e2 = world.spawn(SparseStored(2)).id();

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

@@ -195,7 +195,7 @@ impl SystemExecutor for MultiThreadedExecutor {
             mut conditions,
         } = SyncUnsafeSchedule::new(schedule);
 
-        ComputeTaskPool::init(TaskPool::default).scope_with_executor(
+        ComputeTaskPool::get_or_init(TaskPool::default).scope_with_executor(
             false,
             thread_executor,
             |scope| {

crates/bevy_ecs/src/schedule/mod.rs

@@ -104,7 +104,7 @@ mod tests {
 
             let mut world = World::default();
             let mut schedule = Schedule::default();
-            let thread_count = ComputeTaskPool::init(TaskPool::default).thread_num();
+            let thread_count = ComputeTaskPool::get_or_init(TaskPool::default).thread_num();
 
             let barrier = Arc::new(Barrier::new(thread_count));
 

crates/bevy_tasks/src/usages.rs

@@ -1,107 +1,77 @@
 use super::TaskPool;
 use std::{ops::Deref, sync::OnceLock};
 
-static COMPUTE_TASK_POOL: OnceLock<ComputeTaskPool> = OnceLock::new();
+macro_rules! taskpool {
-static ASYNC_COMPUTE_TASK_POOL: OnceLock<AsyncComputeTaskPool> = OnceLock::new();
+    ($(#[$attr:meta])* ($static:ident, $type:ident)) => {
-static IO_TASK_POOL: OnceLock<IoTaskPool> = OnceLock::new();
+        static $static: OnceLock<$type> = OnceLock::new();
 
-/// A newtype for a task pool for CPU-intensive work that must be completed to
+        $(#[$attr])*
-/// deliver the next frame
+        #[derive(Debug)]
-///
+        pub struct $type(TaskPool);
-/// See [`TaskPool`] documentation for details on Bevy tasks.
-/// [`AsyncComputeTaskPool`] should be preferred if the work does not have to be
-/// completed before the next frame.
-#[derive(Debug)]
-pub struct ComputeTaskPool(TaskPool);
 
-impl ComputeTaskPool {
+        impl $type {
-    /// Initializes the global [`ComputeTaskPool`] instance.
+            #[doc = concat!(" Gets the global [`", stringify!($type), "`] instance, or initializes it with `f`.")]
-    pub fn init(f: impl FnOnce() -> TaskPool) -> &'static Self {
+            pub fn get_or_init(f: impl FnOnce() -> TaskPool) -> &'static Self {
-        COMPUTE_TASK_POOL.get_or_init(|| Self(f()))
+                $static.get_or_init(|| Self(f()))
-    }
+            }
 
-    /// Gets the global [`ComputeTaskPool`] instance.
+            #[doc = concat!(" Attempts to get the global [`", stringify!($type), "`] instance, \
+                or returns `None` if it is not initialized.")]
+            pub fn try_get() -> Option<&'static Self> {
+                $static.get()
+            }
+
+            #[doc = concat!(" Gets the global [`", stringify!($type), "`] instance.")]
+            #[doc = ""]
+            #[doc = " # Panics"]
+            #[doc = " Panics if the global instance has not been initialized yet."]
+            pub fn get() -> &'static Self {
+                $static.get().expect(
+                    concat!(
+                        "The ",
+                        stringify!($type),
+                        " has not been initialized yet. Please call ",
+                        stringify!($type),
+                        "::get_or_init beforehand."
+                    )
+                )
+            }
+        }
+
+        impl Deref for $type {
+            type Target = TaskPool;
+
+            fn deref(&self) -> &Self::Target {
+                &self.0
+            }
+        }
+    };
+}
+
+taskpool! {
+    /// A newtype for a task pool for CPU-intensive work that must be completed to
+    /// deliver the next frame
     ///
-    /// # Panics
+    /// See [`TaskPool`] documentation for details on Bevy tasks.
-    /// Panics if no pool has been initialized yet.
+    /// [`AsyncComputeTaskPool`] should be preferred if the work does not have to be
-    pub fn get() -> &'static Self {
+    /// completed before the next frame.
-        COMPUTE_TASK_POOL.get().expect(
+    (COMPUTE_TASK_POOL, ComputeTaskPool)
-            "A ComputeTaskPool has not been initialized yet. Please call \
-                    ComputeTaskPool::init beforehand.",
-        )
-    }
 }
 
-impl Deref for ComputeTaskPool {
+taskpool! {
-    type Target = TaskPool;
+    /// A newtype for a task pool for CPU-intensive work that may span across multiple frames
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-/// A newtype for a task pool for CPU-intensive work that may span across multiple frames
-///
-/// See [`TaskPool`] documentation for details on Bevy tasks. Use [`ComputeTaskPool`] if
-/// the work must be complete before advancing to the next frame.
-#[derive(Debug)]
-pub struct AsyncComputeTaskPool(TaskPool);
-
-impl AsyncComputeTaskPool {
-    /// Initializes the global [`AsyncComputeTaskPool`] instance.
-    pub fn init(f: impl FnOnce() -> TaskPool) -> &'static Self {
-        ASYNC_COMPUTE_TASK_POOL.get_or_init(|| Self(f()))
-    }
-
-    /// Gets the global [`AsyncComputeTaskPool`] instance.
     ///
-    /// # Panics
+    /// See [`TaskPool`] documentation for details on Bevy tasks.
-    /// Panics if no pool has been initialized yet.
+    /// Use [`ComputeTaskPool`] if the work must be complete before advancing to the next frame.
-    pub fn get() -> &'static Self {
+    (ASYNC_COMPUTE_TASK_POOL, AsyncComputeTaskPool)
-        ASYNC_COMPUTE_TASK_POOL.get().expect(
-            "A AsyncComputeTaskPool has not been initialized yet. Please call \
-                    AsyncComputeTaskPool::init beforehand.",
-        )
-    }
 }
 
-impl Deref for AsyncComputeTaskPool {
+taskpool! {
-    type Target = TaskPool;
+    /// A newtype for a task pool for IO-intensive work (i.e. tasks that spend very little time in a
-
+    /// "woken" state)
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
-}
-
-/// A newtype for a task pool for IO-intensive work (i.e. tasks that spend very little time in a
-/// "woken" state)
-#[derive(Debug)]
-pub struct IoTaskPool(TaskPool);
-
-impl IoTaskPool {
-    /// Initializes the global [`IoTaskPool`] instance.
-    pub fn init(f: impl FnOnce() -> TaskPool) -> &'static Self {
-        IO_TASK_POOL.get_or_init(|| Self(f()))
-    }
-
-    /// Gets the global [`IoTaskPool`] instance.
     ///
-    /// # Panics
+    /// See [`TaskPool`] documentation for details on Bevy tasks.
-    /// Panics if no pool has been initialized yet.
+    (IO_TASK_POOL, IoTaskPool)
-    pub fn get() -> &'static Self {
-        IO_TASK_POOL.get().expect(
-            "A IoTaskPool has not been initialized yet. Please call \
-                    IoTaskPool::init beforehand.",
-        )
-    }
-}
-
-impl Deref for IoTaskPool {
-    type Target = TaskPool;
-
-    fn deref(&self) -> &Self::Target {
-        &self.0
-    }
 }
 
 /// A function used by `bevy_core` to tick the global tasks pools on the main thread.

crates/bevy_transform/src/systems.rs

@@ -193,7 +193,7 @@ mod test {
 
     #[test]
     fn correct_parent_removed() {
-        ComputeTaskPool::init(TaskPool::default);
+        ComputeTaskPool::get_or_init(TaskPool::default);
         let mut world = World::default();
         let offset_global_transform =
             |offset| GlobalTransform::from(Transform::from_xyz(offset, offset, offset));
@@ -248,7 +248,7 @@ mod test {
 
     #[test]
     fn did_propagate() {
-        ComputeTaskPool::init(TaskPool::default);
+        ComputeTaskPool::get_or_init(TaskPool::default);
         let mut world = World::default();
 
         let mut schedule = Schedule::default();
@@ -326,7 +326,7 @@ mod test {
 
     #[test]
     fn correct_children() {
-        ComputeTaskPool::init(TaskPool::default);
+        ComputeTaskPool::get_or_init(TaskPool::default);
         let mut world = World::default();
 
         let mut schedule = Schedule::default();
@@ -404,7 +404,7 @@ mod test {
     #[test]
     fn correct_transforms_when_no_children() {
         let mut app = App::new();
-        ComputeTaskPool::init(TaskPool::default);
+        ComputeTaskPool::get_or_init(TaskPool::default);
 
         app.add_systems(Update, (sync_simple_transforms, propagate_transforms));
 
@@ -446,7 +446,7 @@ mod test {
     #[test]
     #[should_panic]
     fn panic_when_hierarchy_cycle() {
-        ComputeTaskPool::init(TaskPool::default);
+        ComputeTaskPool::get_or_init(TaskPool::default);
         // We cannot directly edit Parent and Children, so we use a temp world to break
         // the hierarchy's invariants.
         let mut temp = World::new();