diff --git a/crates/bevy_tasks/src/task.rs b/crates/bevy_tasks/src/task.rs index eb257080b4207..d495ebce71252 100644 --- a/crates/bevy_tasks/src/task.rs +++ b/crates/bevy_tasks/src/task.rs @@ -1,7 +1,7 @@ use std::{ future::Future, pin::Pin, - task::{Context, Poll}, + task::{Context, Poll, RawWaker, RawWakerVTable, Waker}, }; /// Wraps `async_executor::Task`, a spawned future. @@ -40,6 +40,25 @@ impl Task { pub async fn cancel(self) -> Option { self.0.cancel().await } + + /// Check if the task has been completed, and if so, return the value. + pub fn check(&mut self) -> Option { + const NOOP_WAKER_VTABLE: &RawWakerVTable = &RawWakerVTable::new( + |_| RawWaker::new(std::ptr::null::<()>(), NOOP_WAKER_VTABLE), + |_| {}, + |_| {}, + |_| {}, + ); + + // SAFE: all vtable functions are always safe to call. + let waker = + unsafe { Waker::from_raw(RawWaker::new(std::ptr::null::<()>(), NOOP_WAKER_VTABLE)) }; + + match Pin::new(&mut self.0).poll(&mut Context::from_waker(&waker)) { + Poll::Ready(val) => Some(val), + Poll::Pending => None, + } + } } impl Future for Task { diff --git a/examples/async_tasks/async_compute.rs b/examples/async_tasks/async_compute.rs index 9182ac48797fd..ed886beb57649 100644 --- a/examples/async_tasks/async_compute.rs +++ b/examples/async_tasks/async_compute.rs @@ -2,7 +2,6 @@ use bevy::{ prelude::*, tasks::{AsyncComputeTaskPool, Task}, }; -use futures_lite::future; use rand::Rng; use std::time::{Duration, Instant}; @@ -86,7 +85,7 @@ fn handle_tasks( box_material_handle: Res, ) { for (entity, mut task) in transform_tasks.iter_mut() { - if let Some(transform) = future::block_on(future::poll_once(&mut task.0)) { + if let Some(transform) = task.0.check() { // Add our new PbrBundle of components to our tagged entity commands.entity(entity).insert_bundle(PbrBundle { mesh: box_mesh_handle.clone(),