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`.
This commit is contained in:
Pixelstorm 2023-10-23 21:48:48 +01:00 committed by GitHub
parent 7d504b89c3
commit faa1b57de5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 74 additions and 104 deletions

View file

@ -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();

View file

@ -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())

View file

@ -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();

View file

@ -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| {

View file

@ -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));

View file

@ -1,107 +1,77 @@
use super::TaskPool;
use std::{ops::Deref, sync::OnceLock};
static COMPUTE_TASK_POOL: OnceLock<ComputeTaskPool> = OnceLock::new();
static ASYNC_COMPUTE_TASK_POOL: OnceLock<AsyncComputeTaskPool> = OnceLock::new();
static IO_TASK_POOL: OnceLock<IoTaskPool> = OnceLock::new();
macro_rules! taskpool {
($(#[$attr:meta])* ($static:ident, $type:ident)) => {
static $static: OnceLock<$type> = 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);
$(#[$attr])*
#[derive(Debug)]
pub struct $type(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()))
}
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()))
}
/// 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
/// 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.",
)
}
/// 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)
}
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 may span across multiple frames
///
/// # 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.",
)
}
/// 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 AsyncComputeTaskPool {
type Target = TaskPool;
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.
taskpool! {
/// A newtype for a task pool for IO-intensive work (i.e. tasks that spend very little time in a
/// "woken" state)
///
/// # 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.",
)
}
}
impl Deref for IoTaskPool {
type Target = TaskPool;
fn deref(&self) -> &Self::Target {
&self.0
}
/// 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.

View file

@ -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();