Skip to content

Commit b1293cf

Browse files
bors[bot]bors[bot]
authored
Try #2250:
2 parents 0520947 + f75eb85 commit b1293cf

File tree

5 files changed

+117
-23
lines changed

5 files changed

+117
-23
lines changed

‎crates/bevy_asset/src/lib.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ use bevy_ecs::{
3232
system::IntoSystem,
3333
};
3434
use bevy_tasks::IoTaskPool;
35+
use std::ops::Deref;
3536

3637
/// The names of asset stages in an App Schedule
3738
#[derive(Debug, Hash, PartialEq, Eq, Clone, StageLabel)]
@@ -83,7 +84,7 @@ impl Plugin for AssetPlugin {
8384
.world()
8485
.get_resource::<IoTaskPool>()
8586
.expect("`IoTaskPool` resource not found.")
86-
.0
87+
.deref()
8788
.clone();
8889

8990
let source = create_platform_default_asset_io(app);

‎crates/bevy_core/src/task_pool_options.rs

Lines changed: 31 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -109,12 +109,16 @@ impl DefaultTaskPoolOptions {
109109
trace!("IO Threads: {}", io_threads);
110110
remaining_threads = remaining_threads.saturating_sub(io_threads);
111111

112-
world.insert_resource(IoTaskPool(
113-
TaskPoolBuilder::default()
114-
.num_threads(io_threads)
115-
.thread_name("IO Task Pool".to_string())
116-
.build(),
117-
));
112+
let task_pool = TaskPoolBuilder::default()
113+
.num_threads(io_threads)
114+
.thread_name("IO Task Pool".to_string())
115+
.build();
116+
117+
let io_task_pool = IoTaskPool::init(task_pool)
118+
.map(|pool| pool.clone())
119+
.unwrap_or_else(|_| IoTaskPool::get().clone());
120+
121+
world.insert_resource(io_task_pool);
118122
}
119123

120124
if !world.contains_resource::<AsyncComputeTaskPool>() {
@@ -126,12 +130,16 @@ impl DefaultTaskPoolOptions {
126130
trace!("Async Compute Threads: {}", async_compute_threads);
127131
remaining_threads = remaining_threads.saturating_sub(async_compute_threads);
128132

129-
world.insert_resource(AsyncComputeTaskPool(
130-
TaskPoolBuilder::default()
131-
.num_threads(async_compute_threads)
132-
.thread_name("Async Compute Task Pool".to_string())
133-
.build(),
134-
));
133+
let task_pool = TaskPoolBuilder::default()
134+
.num_threads(async_compute_threads)
135+
.thread_name("Async Compute Task Pool".to_string())
136+
.build();
137+
138+
let async_task_pool = AsyncComputeTaskPool::init(task_pool)
139+
.map(|pool| pool.clone())
140+
.unwrap_or_else(|_| AsyncComputeTaskPool::get().clone());
141+
142+
world.insert_resource(async_task_pool);
135143
}
136144

137145
if !world.contains_resource::<ComputeTaskPool>() {
@@ -142,12 +150,17 @@ impl DefaultTaskPoolOptions {
142150
.get_number_of_threads(remaining_threads, total_threads);
143151

144152
trace!("Compute Threads: {}", compute_threads);
145-
world.insert_resource(ComputeTaskPool(
146-
TaskPoolBuilder::default()
147-
.num_threads(compute_threads)
148-
.thread_name("Compute Task Pool".to_string())
149-
.build(),
150-
));
153+
154+
let task_pool = TaskPoolBuilder::default()
155+
.num_threads(compute_threads)
156+
.thread_name("Compute Task Pool".to_string())
157+
.build();
158+
159+
let compute_task_pool = ComputeTaskPool::init(task_pool)
160+
.map(|pool| pool.clone())
161+
.unwrap_or_else(|_| ComputeTaskPool::get().clone());
162+
163+
world.insert_resource(compute_task_pool);
151164
}
152165
}
153166
}

‎crates/bevy_ecs/src/schedule/executor_parallel.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,11 @@ impl ParallelSystemExecutor for ParallelExecutor {
115115
self.update_archetypes(systems, world);
116116

117117
let compute_pool = world
118-
.get_resource_or_insert_with(|| ComputeTaskPool(TaskPool::default()))
118+
.get_resource_or_insert_with(|| {
119+
ComputeTaskPool::init(TaskPool::default())
120+
.map(|pool| pool.clone())
121+
.unwrap_or_else(|_| ComputeTaskPool::get().clone())
122+
})
119123
.clone();
120124
compute_pool.scope(|scope| {
121125
self.prepare_systems(scope, systems, world);

‎crates/bevy_tasks/Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,5 +22,7 @@ async-executor = "1.3.0"
2222
async-channel = "1.4.2"
2323
instant = { version = "0.1", features = ["wasm-bindgen"] }
2424
num_cpus = "1"
25+
once_cell = "1.7"
26+
2527
[target.'cfg(target_arch = "wasm32")'.dependencies]
2628
wasm-bindgen-futures = "0.4"

‎crates/bevy_tasks/src/usages.rs

Lines changed: 77 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,40 @@
1111
//! for consumption. (likely via channels)
1212
1313
use super::TaskPool;
14+
use once_cell::sync::OnceCell;
1415
use std::ops::Deref;
1516

17+
static COMPUTE_TASK_POOL: OnceCell<ComputeTaskPool> = OnceCell::new();
18+
static ASYNC_COMPUTE_TASK_POOL: OnceCell<AsyncComputeTaskPool> = OnceCell::new();
19+
static IO_TASK_POOL: OnceCell<IoTaskPool> = OnceCell::new();
20+
1621
/// A newtype for a task pool for CPU-intensive work that must be completed to deliver the next
1722
/// frame
1823
#[derive(Clone, Debug)]
19-
pub struct ComputeTaskPool(pub TaskPool);
24+
pub struct ComputeTaskPool(TaskPool);
25+
26+
impl ComputeTaskPool {
27+
/// Initializes the global ComputeTaskPool instance.
28+
///
29+
/// Returns the provided `[TaskPool]` if the global instance has already been initialized.
30+
pub fn init(task_pool: TaskPool) -> Result<&'static Self, TaskPool> {
31+
COMPUTE_TASK_POOL
32+
.set(Self(task_pool))
33+
.map(|_| Self::get())
34+
.map_err(|pool| pool.0)
35+
}
36+
37+
/// Gets the global ComputeTaskPool instance.
38+
///
39+
/// # Panics
40+
/// Panics if no pool has been initialized yet.
41+
pub fn get() -> &'static Self {
42+
COMPUTE_TASK_POOL.get().expect(
43+
"A ComputeTaskPool has not been initialized yet. Please call \
44+
ComputeTaskPool::init beforehand.",
45+
)
46+
}
47+
}
2048

2149
impl Deref for ComputeTaskPool {
2250
type Target = TaskPool;
@@ -28,7 +56,30 @@ impl Deref for ComputeTaskPool {
2856

2957
/// A newtype for a task pool for CPU-intensive work that may span across multiple frames
3058
#[derive(Clone, Debug)]
31-
pub struct AsyncComputeTaskPool(pub TaskPool);
59+
pub struct AsyncComputeTaskPool(TaskPool);
60+
61+
impl AsyncComputeTaskPool {
62+
/// Initializes the global AsyncComputeTaskPool instance.
63+
///
64+
/// Returns the provided `[TaskPool]` if the global instance has already been initialized.
65+
pub fn init(task_pool: TaskPool) -> Result<&'static Self, TaskPool> {
66+
ASYNC_COMPUTE_TASK_POOL
67+
.set(Self(task_pool))
68+
.map(|_| Self::get())
69+
.map_err(|pool| pool.0)
70+
}
71+
72+
/// Gets the global AsyncComputeTaskPool instance.
73+
///
74+
/// # Panics
75+
/// Panics if no pool has been initialized yet.
76+
pub fn get() -> &'static Self {
77+
ASYNC_COMPUTE_TASK_POOL.get().expect(
78+
"A AsyncComputeTaskPool has not been initialized yet. Please call \
79+
AsyncComputeTaskPool::init beforehand.",
80+
)
81+
}
82+
}
3283

3384
impl Deref for AsyncComputeTaskPool {
3485
type Target = TaskPool;
@@ -41,7 +92,30 @@ impl Deref for AsyncComputeTaskPool {
4192
/// A newtype for a task pool for IO-intensive work (i.e. tasks that spend very little time in a
4293
/// "woken" state)
4394
#[derive(Clone, Debug)]
44-
pub struct IoTaskPool(pub TaskPool);
95+
pub struct IoTaskPool(TaskPool);
96+
97+
impl IoTaskPool {
98+
/// Initializes the global IoTaskPool instance.
99+
///
100+
/// Returns the provided `[TaskPool]` if the global instance has already been initialized.
101+
pub fn init(task_pool: TaskPool) -> Result<&'static Self, TaskPool> {
102+
IO_TASK_POOL
103+
.set(Self(task_pool))
104+
.map(|_| Self::get())
105+
.map_err(|pool| pool.0)
106+
}
107+
108+
/// Gets the global IoTaskPool instance.
109+
///
110+
/// # Panics
111+
/// Panics if no pool has been initialized yet.
112+
pub fn get() -> &'static Self {
113+
IO_TASK_POOL.get().expect(
114+
"A IoTaskPool has not been initialized yet. Please call \
115+
IoTaskPool::init beforehand.",
116+
)
117+
}
118+
}
45119

46120
impl Deref for IoTaskPool {
47121
type Target = TaskPool;

0 commit comments

Comments
 (0)