diff --git a/crates/guest-rust/src/rt/async_support.rs b/crates/guest-rust/src/rt/async_support.rs index d2164bdf2..e29c5a234 100644 --- a/crates/guest-rust/src/rt/async_support.rs +++ b/crates/guest-rust/src/rt/async_support.rs @@ -235,15 +235,11 @@ impl FutureState<'_> { let poll = me.tasks.poll_next(&mut context); match poll { - // A future completed, yay! Keep going to see if more have - // completed. - Poll::Ready(Some(())) => (), - // The task list is empty, but there might be remaining work // in terms of waitables through the cabi interface. In this // situation wait for all waitables to be resolved before // signaling that our own task is done. - Poll::Ready(None) => { + Poll::Ready(()) => { assert!(me.tasks.is_empty()); if me.remaining_work() { let waitable = me.waitable_set.as_ref().unwrap().as_raw(); diff --git a/crates/guest-rust/src/rt/async_support/spawn.rs b/crates/guest-rust/src/rt/async_support/spawn.rs index 2a5c5659e..083f686e4 100644 --- a/crates/guest-rust/src/rt/async_support/spawn.rs +++ b/crates/guest-rust/src/rt/async_support/spawn.rs @@ -26,13 +26,46 @@ impl<'a> Tasks<'a> { } } - pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll> { - unsafe { - let ret = self.tasks.poll_next_unpin(cx); - if !SPAWNED.is_empty() { - self.tasks.extend(SPAWNED.drain(..)); + pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<()> { + loop { + // Perform some work by seeing what's next in this + // `FuturesUnordered`. Afterwards check the set of spawned tasks + // and, if any, add them to our set of tasks being done. + let poll = self.tasks.poll_next_unpin(cx); + let spawned = unsafe { + if SPAWNED.is_empty() { + false + } else { + self.tasks.extend(SPAWNED.drain(..)); + true + } + }; + match poll { + // If no tasks were ready, and if we didn't spawn any work, + // then there's nothing left to do so return pending. + // + // If no tasks were ready, and if we spawned some work, then + // turn the loop again to register interest in the work and + // ensure that it's not forgotten about. + Poll::Pending => { + if !spawned { + return Poll::Pending; + } + } + + // If our set of tasks is empty it shouldn't be possible to have + // spawned anything, and return saying that we're done. + Poll::Ready(None) => { + assert!(!spawned); + return Poll::Ready(()); + } + + // If a task finished, then turn the loop again to see if there + // are any other completed tasks. This also serves double-duty + // to ensure that we look at everything in our set of tasks + // before concluding that we're finished. + Poll::Ready(Some(())) => {} } - ret } } diff --git a/crates/guest-rust/src/rt/async_support/spawn_disabled.rs b/crates/guest-rust/src/rt/async_support/spawn_disabled.rs index 7840cf3b2..414b746af 100644 --- a/crates/guest-rust/src/rt/async_support/spawn_disabled.rs +++ b/crates/guest-rust/src/rt/async_support/spawn_disabled.rs @@ -11,14 +11,14 @@ impl<'a> Tasks<'a> { Tasks { future: Some(root) } } - pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll> { + pub fn poll_next(&mut self, cx: &mut Context<'_>) -> Poll<()> { if let Some(future) = self.future.as_mut() { if future.as_mut().poll(cx).is_ready() { self.future = None; } } if self.is_empty() { - Poll::Ready(None) + Poll::Ready(()) } else { Poll::Pending }