diff --git a/benches/benches/bevy_ecs/iteration/heavy_compute.rs b/benches/benches/bevy_ecs/iteration/heavy_compute.rs index c1b8598b97e972..9a53092903f489 100644 --- a/benches/benches/bevy_ecs/iteration/heavy_compute.rs +++ b/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(); diff --git a/crates/bevy_core/src/task_pool_options.rs b/crates/bevy_core/src/task_pool_options.rs index 3eaedb8c8f9723..a6eb39df6c7236 100644 --- a/crates/bevy_core/src/task_pool_options.rs +++ b/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()) diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs index f502c7fb34c3b9..71109585f2c57b 100644 --- a/crates/bevy_ecs/src/lib.rs +++ b/crates/bevy_ecs/src/lib.rs @@ -399,7 +399,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(); @@ -422,7 +422,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(); diff --git a/crates/bevy_ecs/src/schedule/executor/multi_threaded.rs b/crates/bevy_ecs/src/schedule/executor/multi_threaded.rs index dad3fa4459d20f..36eb76a94abcaa 100644 --- a/crates/bevy_ecs/src/schedule/executor/multi_threaded.rs +++ b/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| { diff --git a/crates/bevy_ecs/src/schedule/mod.rs b/crates/bevy_ecs/src/schedule/mod.rs index 4ed863944acfb0..9d48d03df40e90 100644 --- a/crates/bevy_ecs/src/schedule/mod.rs +++ b/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)); diff --git a/crates/bevy_tasks/src/usages.rs b/crates/bevy_tasks/src/usages.rs index 49b8b5cd2ff729..fda3092b8ebc85 100644 --- a/crates/bevy_tasks/src/usages.rs +++ b/crates/bevy_tasks/src/usages.rs @@ -1,107 +1,77 @@ use super::TaskPool; use std::{ops::Deref, sync::OnceLock}; -static COMPUTE_TASK_POOL: OnceLock = OnceLock::new(); -static ASYNC_COMPUTE_TASK_POOL: OnceLock = OnceLock::new(); -static IO_TASK_POOL: OnceLock = OnceLock::new(); - -/// A newtype for a task pool for CPU-intensive work that must be completed to -/// deliver the next frame -/// -/// 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 { - /// Initializes the global [`ComputeTaskPool`] instance. - pub fn init(f: impl FnOnce() -> TaskPool) -> &'static Self { - COMPUTE_TASK_POOL.get_or_init(|| Self(f())) - } - - /// Gets the global [`ComputeTaskPool`] instance. - /// - /// # Panics - /// Panics if no pool has been initialized yet. - pub fn get() -> &'static Self { - COMPUTE_TASK_POOL.get().expect( - "A ComputeTaskPool has not been initialized yet. Please call \ - ComputeTaskPool::init beforehand.", - ) - } +macro_rules! taskpool { + ($(#[$attr:meta])* ($static:ident, $type:ident)) => { + static $static: OnceLock<$type> = OnceLock::new(); + + $(#[$attr])* + #[derive(Debug)] + pub struct $type(TaskPool); + + impl $type { + #[doc = concat!(" Gets the global [`", stringify!($type), "`] instance, or initializes it with `f`.")] + pub fn get_or_init(f: impl FnOnce() -> TaskPool) -> &'static Self { + $static.get_or_init(|| Self(f())) + } + + #[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 + } + } + }; } -impl Deref for ComputeTaskPool { - type Target = TaskPool; - - 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. +taskpool! { + /// A newtype for a task pool for CPU-intensive work that must be completed to + /// deliver the next frame /// - /// # Panics - /// Panics if no pool has been initialized yet. - pub fn get() -> &'static Self { - ASYNC_COMPUTE_TASK_POOL.get().expect( - "A AsyncComputeTaskPool has not been initialized yet. Please call \ - AsyncComputeTaskPool::init beforehand.", - ) - } -} - -impl Deref for AsyncComputeTaskPool { - type Target = TaskPool; - - fn deref(&self) -> &Self::Target { - &self.0 - } + /// 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. + (COMPUTE_TASK_POOL, ComputeTaskPool) } -/// 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. +taskpool! { + /// A newtype for a task pool for CPU-intensive work that may span across multiple frames /// - /// # Panics - /// Panics if no pool has been initialized yet. - pub fn get() -> &'static Self { - IO_TASK_POOL.get().expect( - "A IoTaskPool has not been initialized yet. Please call \ - IoTaskPool::init beforehand.", - ) - } + /// See [`TaskPool`] documentation for details on Bevy tasks. + /// Use [`ComputeTaskPool`] if the work must be complete before advancing to the next frame. + (ASYNC_COMPUTE_TASK_POOL, AsyncComputeTaskPool) } -impl Deref for IoTaskPool { - type Target = TaskPool; - - fn deref(&self) -> &Self::Target { - &self.0 - } +taskpool! { + /// A newtype for a task pool for IO-intensive work (i.e. tasks that spend very little time in a + /// "woken" state) + /// + /// See [`TaskPool`] documentation for details on Bevy tasks. + (IO_TASK_POOL, IoTaskPool) } /// A function used by `bevy_core` to tick the global tasks pools on the main thread. diff --git a/crates/bevy_transform/src/systems.rs b/crates/bevy_transform/src/systems.rs index b25784c3889fe4..8f6bac916a7395 100644 --- a/crates/bevy_transform/src/systems.rs +++ b/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();