From b2732669af450e67dadfc059659b348808157d9e Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Tue, 28 Apr 2026 13:02:29 -0700
Subject: [PATCH] Adjust spawning behavior and when tasks are polled

This commit adjusts the behavior of the `async-spawn` feature by notably
ensuring that spawned tasks are polled even when the main task is
complete. This ensures that spawned work is looked at as opposed to
being stuck in limbo (never polled). This resolves an issue with
bytecodealliance/wasmtime#13196, for example, where work was spawned,
but the main set of work was complete, so the specific interleaving
there meant that the spawned work never actually ended up happening.
---
 crates/guest-rust/src/rt/async_support.rs     |  6 +--
 .../guest-rust/src/rt/async_support/spawn.rs  | 45 ++++++++++++++++---
 .../src/rt/async_support/spawn_disabled.rs    |  4 +-
 3 files changed, 42 insertions(+), 13 deletions(-)

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<Option<()>> {
-        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<Option<()>> {
+    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
         }