diff --git a/crates/bevy_tasks/src/task_pool.rs b/crates/bevy_tasks/src/task_pool.rs index 4f82e02241203..74ba7c92e3464 100644 --- a/crates/bevy_tasks/src/task_pool.rs +++ b/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> + Send)> = fut; + let fut: Pin<&'static mut (dyn Future> + 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> + Send)> = fut; - let fut: Pin<&'static mut (dyn Future> + 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(); } }