When a task scope produces <= 1 task to run, run it on the calling thread immediately. (#932)

While generally speaking the calling thread would have picked up the task first anyways, I don't think it makes much sense usually to block the calling thread until another thread wakes and does the work.

Author: aclysma

crates/bevy_tasks/src/task_pool.rs

@@ -167,44 +167,50 @@ impl TaskPool {
         let executor: &async_executor::Executor = &*self.executor;
         let executor: &'scope async_executor::Executor = unsafe { mem::transmute(executor) };
 
-        let fut = async move {
-            let mut scope = Scope {
-                executor,
-                spawned: Vec::new(),
-            };
-
-            f(&mut scope);
+        let mut scope = Scope {
+            executor,
+            spawned: Vec::new(),
+        };
 
-            let mut results = Vec::with_capacity(scope.spawned.len());
-            for task in scope.spawned {
-                results.push(task.await);
-            }
+        f(&mut scope);
+
+        if scope.spawned.is_empty() {
+            Vec::default()
+        } else if scope.spawned.len() == 1 {
+            vec![future::block_on(&mut scope.spawned[0])]
+        } else {
+            let fut = async move {
+                let mut results = Vec::with_capacity(scope.spawned.len());
+                for task in scope.spawned {
+                    results.push(task.await);
+                }
 
-            results
-        };
+                results
+            };
 
-        // Pin the future on the stack.
-        pin!(fut);
+            // Pin the future on the stack.
+            pin!(fut);
+
+            // SAFETY: This function blocks until all futures complete, so we do not read/write the
+            // data from futures outside of the 'scope lifetime. However, rust has no way of knowing
+            // this so we must convert to 'static here to appease the compiler as it is unable to
+            // validate safety.
+            let fut: Pin<&mut (dyn Future<Output = Vec<T>> + Send)> = fut;
+            let fut: Pin<&'static mut (dyn Future<Output = Vec<T>> + Send + 'static)> =
+                unsafe { mem::transmute(fut) };
+
+            // The thread that calls scope() will participate in driving tasks in the pool forward
+            // until the tasks that are spawned by this scope() call complete. (If the caller of scope()
+            // happens to be a thread in this thread pool, and we only have one thread in the pool, then
+            // simply calling future::block_on(spawned) would deadlock.)
+            let mut spawned = self.executor.spawn(fut);
+            loop {
+                if let Some(result) = future::block_on(future::poll_once(&mut spawned)) {
+                    break result;
+                }
 
-        // SAFETY: This function blocks until all futures complete, so we do not read/write the
-        // data from futures outside of the 'scope lifetime. However, rust has no way of knowing
-        // this so we must convert to 'static here to appease the compiler as it is unable to
-        // validate safety.
-        let fut: Pin<&mut (dyn Future<Output = Vec<T>> + Send)> = fut;
-        let fut: Pin<&'static mut (dyn Future<Output = Vec<T>> + Send + 'static)> =
-            unsafe { mem::transmute(fut) };
-
-        // The thread that calls scope() will participate in driving tasks in the pool forward
-        // until the tasks that are spawned by this scope() call complete. (If the caller of scope()
-        // happens to be a thread in this thread pool, and we only have one thread in the pool, then
-        // simply calling future::block_on(spawned) would deadlock.)
-        let mut spawned = self.executor.spawn(fut);
-        loop {
-            if let Some(result) = future::block_on(future::poll_once(&mut spawned)) {
-                break result;
+                self.executor.try_tick();
             }
-
-            self.executor.try_tick();
         }
     }