[tasks] Handle TaskPool panicking threads

crates/bevy_core/src/lib.rs

@@ -46,6 +46,24 @@ impl Plugin for CorePlugin {
             .unwrap_or_else(DefaultTaskPoolOptions::default)
             .create_default_pools(app.world_mut());
 
+        #[cfg(not(target_arch = "wasm32"))]
+        app.add_system(
+            task_pool_options::handle_task_pool_panicking_threads_system::<bevy_tasks::IoTaskPool>
+                .system(),
+        )
+        .add_system(
+            task_pool_options::handle_task_pool_panicking_threads_system::<
+                bevy_tasks::ComputeTaskPool,
+            >
+                .system(),
+        )
+        .add_system(
+            task_pool_options::handle_task_pool_panicking_threads_system::<
+                bevy_tasks::AsyncComputeTaskPool,
+            >
+                .system(),
+        );
+
         app.init_resource::<Time>()
             .init_resource::<EntityLabels>()
             .init_resource::<FixedTimesteps>()

crates/bevy_core/src/task_pool_options.rs

@@ -1,10 +1,24 @@
+#[cfg(not(target_arch = "wasm32"))]
+use bevy_ecs::system::Res;
 use bevy_ecs::world::World;
-use bevy_tasks::{AsyncComputeTaskPool, ComputeTaskPool, IoTaskPool, TaskPoolBuilder};
+use bevy_tasks::{
+    AsyncComputeTaskPool, ComputeTaskPool, DerefTaskPool, IoTaskPool, TaskPoolBuilder,
+    TaskPoolThreadPanicPolicy,
+};
 use bevy_utils::tracing::trace;
 
+/// The set of policies describing how the according task pool behaves
+#[derive(Clone, Debug)]
+pub struct TaskPoolPolicies {
+    /// Used to determine number of threads to allocate
+    pub assignment_policy: TaskPoolThreadAssignmentPolicy,
+    /// Used to determine the panic policy of the task pool
+    pub panic_policy: TaskPoolThreadPanicPolicy,
+}
+
 /// Defines a simple way to determine how many threads to use given the number of remaining cores
 /// and number of total cores
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct TaskPoolThreadAssignmentPolicy {
     /// Force using at least this many threads
     pub min_threads: usize,
@@ -42,13 +56,12 @@ pub struct DefaultTaskPoolOptions {
     /// If the number of physical cores is grater than max_total_threads, force using
     /// max_total_threads
     pub max_total_threads: usize,
-
-    /// Used to determine number of IO threads to allocate
-    pub io: TaskPoolThreadAssignmentPolicy,
-    /// Used to determine number of async compute threads to allocate
-    pub async_compute: TaskPoolThreadAssignmentPolicy,
-    /// Used to determine number of compute threads to allocate
-    pub compute: TaskPoolThreadAssignmentPolicy,
+    /// Used to configure the IOTaskPool's inner policies
+    pub io: TaskPoolPolicies,
+    /// Used to configure the AsyncTaskPool's inner policies
+    pub async_compute: TaskPoolPolicies,
+    /// Used to configure the ComputeTaskPool's inner policies
+    pub compute: TaskPoolPolicies,
 }
 
 impl Default for DefaultTaskPoolOptions {
@@ -59,24 +72,33 @@ impl Default for DefaultTaskPoolOptions {
             max_total_threads: std::usize::MAX,
 
             // Use 25% of cores for IO, at least 1, no more than 4
-            io: TaskPoolThreadAssignmentPolicy {
-                min_threads: 1,
-                max_threads: 4,
-                percent: 0.25,
+            io: TaskPoolPolicies {
+                assignment_policy: TaskPoolThreadAssignmentPolicy {
+                    min_threads: 1,
+                    max_threads: 4,
+                    percent: 0.25,
+                },
+                panic_policy: TaskPoolThreadPanicPolicy::Restart,
             },
 
             // Use 25% of cores for async compute, at least 1, no more than 4
-            async_compute: TaskPoolThreadAssignmentPolicy {
-                min_threads: 1,
-                max_threads: 4,
-                percent: 0.25,
+            async_compute: TaskPoolPolicies {
+                assignment_policy: TaskPoolThreadAssignmentPolicy {
+                    min_threads: 1,
+                    max_threads: 4,
+                    percent: 0.25,
+                },
+                panic_policy: TaskPoolThreadPanicPolicy::Restart,
             },
 
             // Use all remaining cores for compute (at least 1)
-            compute: TaskPoolThreadAssignmentPolicy {
-                min_threads: 1,
-                max_threads: std::usize::MAX,
-                percent: 1.0, // This 1.0 here means "whatever is left over"
+            compute: TaskPoolPolicies {
+                assignment_policy: TaskPoolThreadAssignmentPolicy {
+                    min_threads: 1,
+                    max_threads: std::usize::MAX,
+                    percent: 1.0, // This 1.0 here means "whatever is left over"
+                },
+                panic_policy: TaskPoolThreadPanicPolicy::Restart,
             },
         }
     }
@@ -104,6 +126,7 @@ impl DefaultTaskPoolOptions {
             // Determine the number of IO threads we will use
             let io_threads = self
                 .io
+                .assignment_policy
                 .get_number_of_threads(remaining_threads, total_threads);
 
             trace!("IO Threads: {}", io_threads);
@@ -113,6 +136,7 @@ impl DefaultTaskPoolOptions {
                 TaskPoolBuilder::default()
                     .num_threads(io_threads)
                     .thread_name("IO Task Pool".to_string())
+                    .panic_policy(self.io.panic_policy)
                     .build(),
             ));
         }
@@ -121,6 +145,7 @@ impl DefaultTaskPoolOptions {
             // Determine the number of async compute threads we will use
             let async_compute_threads = self
                 .async_compute
+                .assignment_policy
                 .get_number_of_threads(remaining_threads, total_threads);
 
             trace!("Async Compute Threads: {}", async_compute_threads);
@@ -130,6 +155,7 @@ impl DefaultTaskPoolOptions {
                 TaskPoolBuilder::default()
                     .num_threads(async_compute_threads)
                     .thread_name("Async Compute Task Pool".to_string())
+                    .panic_policy(self.async_compute.panic_policy)
                     .build(),
             ));
         }
@@ -139,15 +165,26 @@ impl DefaultTaskPoolOptions {
             // This is intentionally last so that an end user can specify 1.0 as the percent
             let compute_threads = self
                 .compute
+                .assignment_policy
                 .get_number_of_threads(remaining_threads, total_threads);
 
             trace!("Compute Threads: {}", compute_threads);
             world.insert_resource(ComputeTaskPool(
                 TaskPoolBuilder::default()
                     .num_threads(compute_threads)
                     .thread_name("Compute Task Pool".to_string())
+                    .panic_policy(self.compute.panic_policy)
                     .build(),
             ));
         }
     }
 }
+
+#[cfg(not(target_arch = "wasm32"))]
+pub(crate) fn handle_task_pool_panicking_threads_system<
+    T: DerefTaskPool + Send + Sync + 'static,
+>(
+    task_pool: Res<T>,
+) {
+    bevy_tasks::handle_task_pool_panicking_threads(&**task_pool);
+}

crates/bevy_tasks/Cargo.toml

@@ -16,6 +16,8 @@ keywords = ["bevy"]
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+bevy_utils = { path = "../bevy_utils", version = "0.5.0" }
+parking_lot = "0.11.0"
 futures-lite = "1.4.0"
 event-listener = "2.4.0"
 async-executor = "1.3.0"

crates/bevy_tasks/src/lib.rs

@@ -4,18 +4,21 @@ pub use slice::{ParallelSlice, ParallelSliceMut};
 mod task;
 pub use task::Task;
 
+mod task_pool_common;
+pub use task_pool_common::TaskPoolThreadPanicPolicy;
+
 #[cfg(not(target_arch = "wasm32"))]
 mod task_pool;
 #[cfg(not(target_arch = "wasm32"))]
-pub use task_pool::{Scope, TaskPool, TaskPoolBuilder};
+pub use task_pool::{handle_task_pool_panicking_threads, Scope, TaskPool, TaskPoolBuilder};
 
 #[cfg(target_arch = "wasm32")]
 mod single_threaded_task_pool;
 #[cfg(target_arch = "wasm32")]
 pub use single_threaded_task_pool::{Scope, TaskPool, TaskPoolBuilder};
 
 mod usages;
-pub use usages::{AsyncComputeTaskPool, ComputeTaskPool, IoTaskPool};
+pub use usages::{AsyncComputeTaskPool, ComputeTaskPool, DerefTaskPool, IoTaskPool};
 
 mod countdown_event;
 pub use countdown_event::CountdownEvent;

crates/bevy_tasks/src/single_threaded_task_pool.rs

@@ -1,3 +1,4 @@
+use crate::TaskPoolThreadPanicPolicy;
 use std::{
     future::Future,
     mem,
@@ -26,6 +27,10 @@ impl TaskPoolBuilder {
         self
     }
 
+    pub fn panic_policy(self, _policy: TaskPoolThreadPanicPolicy) -> Self {
+        self
+    }
+
     pub fn build(self) -> TaskPool {
         TaskPool::new_internal()
     }