Add try_spawn and force_spawn scoped tasks methods

Author: ictrobot

crates/utils/src/multithreading/impl_native.rs

@@ -22,8 +22,7 @@ pub fn get_thread_count() -> NonZeroUsize {
 
 /// Set the number of worker threads to use.
 ///
-/// This will affect any future call to [`get_thread_count`], unless scoped tasks are enabled and
-/// the thread pool has already been created.
+/// This will affect any future call to [`get_thread_count`].
 pub fn set_thread_count(count: NonZeroUsize) {
     NUM_THREADS.store(count.get(), Relaxed);
 }
@@ -40,31 +39,21 @@ pub fn worker_pool(worker: impl Fn() + Copy + Send) {
     } else {
         #[cfg(feature = "scoped-tasks")]
         {
-            use super::scoped_tasks;
-
-            static ONCE: std::sync::Once = std::sync::Once::new();
-            ONCE.call_once(|| {
-                for i in 0..threads {
-                    std::thread::Builder::new()
-                        .name(format!("worker-{i}"))
-                        .spawn(scoped_tasks::worker)
-                        .expect("failed to spawn worker thread");
-                }
-            });
-
-            scoped_tasks::scope(|scope| {
-                for _ in 0..threads {
-                    scope.spawn(worker);
+            super::scoped_tasks::scope(|scope| {
+                for _ in 1..threads {
+                    scope.force_spawn(worker);
                 }
+                worker();
             });
         }
 
         #[cfg(not(feature = "scoped-tasks"))]
         {
-            std::thread::scope(|scope| {
-                for _ in 0..threads {
+            std::thread::scope(move |scope| {
+                for _ in 1..threads {
                     scope.spawn(worker);
                 }
+                worker();
             });
         }
     }

crates/utils/src/multithreading/scoped_tasks.rs

@@ -69,6 +69,7 @@ use std::any::Any;
 use std::collections::VecDeque;
 use std::marker::PhantomData;
 use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::sync::mpsc::{SyncSender, TrySendError};
 use std::sync::{Arc, Condvar, Mutex};
 
@@ -111,7 +112,7 @@ where
 
 /// Scope to spawn tasks in.
 ///
-/// Designed to match the [`std::thread::Scope`] API.
+/// [`Scope::spawn()`] is designed to match the [`std::thread::Scope`] API.
 ///
 /// # Lifetimes
 ///
@@ -147,6 +148,55 @@ impl<'scope> Scope<'scope, '_> {
         handle
     }
 
+    /// Spawn a new task within the scope if there is a worker available.
+    ///
+    /// If no workers within the thread pool are available, the task will not be executed and
+    /// [`None`] will be returned.
+    pub fn try_spawn<F, T>(&'scope self, f: F) -> Option<ScopedJoinHandle<'scope, T>>
+    where
+        F: FnOnce() -> T + Send + 'scope,
+        T: Send + 'scope,
+    {
+        let (closure, handle) = self.create_closure(f);
+        if let Ok(()) = try_queue_task(closure) {
+            Some(handle)
+        } else {
+            // Closure will never be run
+            self.data.task_end();
+
+            None
+        }
+    }
+
+    /// Spawn a new task within the scope, spawning a new worker if necessary.
+    ///
+    /// This function is not available on WebAssembly, as new threads have to be created from the
+    /// host JS environment.
+    #[cfg(not(target_family = "wasm"))]
+    pub fn force_spawn<F, T>(&'scope self, f: F) -> ScopedJoinHandle<'scope, T>
+    where
+        F: FnOnce() -> T + Send + 'scope,
+        T: Send + 'scope,
+    {
+        let (closure, handle) = self.create_closure(f);
+        if let Err(closure) = try_queue_task(closure) {
+            // Start a worker to process this closure and then join the pool.
+            static THREAD_NUM: AtomicU32 = AtomicU32::new(1);
+            std::thread::Builder::new()
+                .name(format!(
+                    "scoped-tasks-{}",
+                    THREAD_NUM.fetch_add(1, Ordering::Relaxed)
+                ))
+                .spawn(move || {
+                    // Pass the closure directly to the new worker to avoid race conditions where
+                    // another scope queues a closure before this one.
+                    worker_impl(closure);
+                })
+                .expect("failed to spawn worker thread");
+        }
+        handle
+    }
+
     fn create_closure<F, T>(
         &'scope self,
         f: F,
@@ -345,6 +395,12 @@ fn try_queue_task(mut closure: Box<dyn FnOnce() + Send>) -> Result<(), Box<dyn F
 ///
 /// This function never returns.
 pub fn worker() {
+    worker_impl(|| {});
+}
+
+fn worker_impl(initial: impl FnOnce() + Send) {
+    initial();
+
     let (tx, rx) = std::sync::mpsc::sync_channel(0);
 
     {