make thread.yield a no-op in non-blocking contexts

Per the proposed spec changes, `thread.yield` should return control to the guest
immediately without allowing any other thread to run.  Similarly, when an
async-lifted export or callback returns `CALLBACK_CODE_YIELD`, we should call
the callback again immediately without allowing another thread to run.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

Author: dicej

crates/wasmtime/src/runtime/component/concurrent.rs

@@ -600,7 +600,10 @@ enum GuestCallKind {
     },
     /// Indicates that a new guest task call is pending and may be executed
     /// using the specified closure.
-    StartImplicit(Box<dyn FnOnce(&mut dyn VMStore) -> Result<()> + Send + Sync>),
+    ///
+    /// If the closure returns `Ok(Some(call))`, the `call` should be run
+    /// immediately using `handle_guest_call`.
+    StartImplicit(Box<dyn FnOnce(&mut dyn VMStore) -> Result<Option<GuestCall>> + Send + Sync>),
     StartExplicit(Box<dyn FnOnce(&mut dyn VMStore) -> Result<()> + Send + Sync>),
 }
 
@@ -818,53 +821,58 @@ pub(crate) fn poll_and_block<R: Send + Sync + 'static>(
 
 /// Execute the specified guest call.
 fn handle_guest_call(store: &mut dyn VMStore, call: GuestCall) -> Result<()> {
-    match call.kind {
-        GuestCallKind::DeliverEvent { instance, set } => {
-            let (event, waitable) = instance
-                .get_event(store, call.thread.task, set, true)?
-                .unwrap();
-            let state = store.concurrent_state_mut();
-            let task = state.get_mut(call.thread.task)?;
-            let runtime_instance = task.instance;
-            let handle = waitable.map(|(_, v)| v).unwrap_or(0);
+    let mut next = Some(call);
+    while let Some(call) = next.take() {
+        match call.kind {
+            GuestCallKind::DeliverEvent { instance, set } => {
+                let (event, waitable) = instance
+                    .get_event(store, call.thread.task, set, true)?
+                    .unwrap();
+                let state = store.concurrent_state_mut();
+                let task = state.get_mut(call.thread.task)?;
+                let runtime_instance = task.instance;
+                let handle = waitable.map(|(_, v)| v).unwrap_or(0);
 
-            log::trace!(
-                "use callback to deliver event {event:?} to {:?} for {waitable:?}",
-                call.thread,
-            );
+                log::trace!(
+                    "use callback to deliver event {event:?} to {:?} for {waitable:?}",
+                    call.thread,
+                );
 
-            let old_thread = state.guest_thread.replace(call.thread);
-            log::trace!(
-                "GuestCallKind::DeliverEvent: replaced {old_thread:?} with {:?} as current thread",
-                call.thread
-            );
+                let old_thread = state.guest_thread.replace(call.thread);
+                log::trace!(
+                    "GuestCallKind::DeliverEvent: replaced {old_thread:?} with {:?} as current thread",
+                    call.thread
+                );
 
-            store.maybe_push_call_context(call.thread.task)?;
+                store.maybe_push_call_context(call.thread.task)?;
 
-            let state = store.concurrent_state_mut();
-            state.enter_instance(runtime_instance);
+                let state = store.concurrent_state_mut();
+                state.enter_instance(runtime_instance);
 
-            let callback = state.get_mut(call.thread.task)?.callback.take().unwrap();
+                let callback = state.get_mut(call.thread.task)?.callback.take().unwrap();
 
-            let code = callback(store, runtime_instance, event, handle)?;
+                let code = callback(store, runtime_instance, event, handle)?;
 
-            let state = store.concurrent_state_mut();
+                let state = store.concurrent_state_mut();
 
-            state.get_mut(call.thread.task)?.callback = Some(callback);
-            state.exit_instance(runtime_instance)?;
+                state.get_mut(call.thread.task)?.callback = Some(callback);
+                state.exit_instance(runtime_instance)?;
 
-            store.maybe_pop_call_context(call.thread.task)?;
+                store.maybe_pop_call_context(call.thread.task)?;
 
-            instance.handle_callback_code(store, call.thread, runtime_instance, code)?;
+                next = instance.handle_callback_code(store, call.thread, runtime_instance, code)?;
 
-            store.concurrent_state_mut().guest_thread = old_thread;
-            log::trace!("GuestCallKind::DeliverEvent: restored {old_thread:?} as current thread");
-        }
-        GuestCallKind::StartImplicit(fun) => {
-            fun(store)?;
-        }
-        GuestCallKind::StartExplicit(fun) => {
-            fun(store)?;
+                store.concurrent_state_mut().guest_thread = old_thread;
+                log::trace!(
+                    "GuestCallKind::DeliverEvent: restored {old_thread:?} as current thread"
+                );
+            }
+            GuestCallKind::StartImplicit(fun) => {
+                next = fun(store)?;
+            }
+            GuestCallKind::StartExplicit(fun) => {
+                fun(store)?;
+            }
         }
     }
 
@@ -1589,13 +1597,16 @@ impl Instance {
 
     /// Handle the `CallbackCode` returned from an async-lifted export or its
     /// callback.
+    ///
+    /// If this returns `Ok(Some(call))`, then `call` should be run immediately
+    /// using `handle_guest_call`.
     fn handle_callback_code(
         self,
         store: &mut StoreOpaque,
         guest_thread: QualifiedThreadId,
         runtime_instance: RuntimeComponentInstanceIndex,
         code: u32,
-    ) -> Result<()> {
+    ) -> Result<Option<GuestCall>> {
         let (code, set) = unpack_callback_code(code);
 
         log::trace!("received callback code from {guest_thread:?}: {code} (set: {set})");
@@ -1613,7 +1624,7 @@ impl Instance {
             Ok(TableId::<WaitableSet>::new(set))
         };
 
-        match code {
+        Ok(match code {
             callback_code::EXIT => {
                 log::trace!("implicit thread {guest_thread:?} completed");
                 self.cleanup_thread(store, guest_thread, runtime_instance)?;
@@ -1633,20 +1644,30 @@ impl Instance {
                         task.callback = None;
                     }
                 }
+                None
             }
             callback_code::YIELD => {
-                // Push this thread onto the "low priority" queue so it runs after
-                // any other threads have had a chance to run.
                 let task = state.get_mut(guest_thread.task)?;
                 assert!(task.event.is_none());
                 task.event = Some(Event::None);
-                state.push_low_priority(WorkItem::GuestCall(GuestCall {
+                let call = GuestCall {
                     thread: guest_thread,
                     kind: GuestCallKind::DeliverEvent {
                         instance: self,
                         set: None,
                     },
-                }));
+                };
+                if state.may_block(guest_thread.task) {
+                    // Push this thread onto the "low priority" queue so it runs
+                    // after any other threads have had a chance to run.
+                    state.push_low_priority(WorkItem::GuestCall(call));
+                    None
+                } else {
+                    // Yielding in a non-blocking context is defined as a no-op
+                    // according to the spec, so we must run this thread
+                    // immediately without allowing any others to run.
+                    Some(call)
+                }
             }
             callback_code::WAIT | callback_code::POLL => {
                 state.check_blocking_for(guest_thread.task)?;
@@ -1698,11 +1719,10 @@ impl Instance {
                         _ => unreachable!(),
                     }
                 }
+                None
             }
             _ => bail!("unsupported callback code: {code}"),
-        }
-
-        Ok(())
+        })
     }
 
     fn cleanup_thread(
@@ -1872,10 +1892,9 @@ impl Instance {
                 // function returns a `i32` result.
                 let code = unsafe { storage[0].assume_init() }.get_i32() as u32;
 
-                self.handle_callback_code(store, guest_thread, callee_instance, code)?;
-
-                Ok(())
-            }) as Box<dyn FnOnce(&mut dyn VMStore) -> Result<()> + Send + Sync>
+                self.handle_callback_code(store, guest_thread, callee_instance, code)
+            })
+                as Box<dyn FnOnce(&mut dyn VMStore) -> Result<Option<GuestCall>> + Send + Sync>
         } else {
             let token = StoreToken::new(store.as_context_mut());
             Box::new(move |store: &mut dyn VMStore| {
@@ -2011,7 +2030,7 @@ impl Instance {
                     }
                 }
 
-                Ok(())
+                Ok(None)
             })
         };
 
@@ -3079,9 +3098,20 @@ impl Instance {
     ) -> Result<WaitResult> {
         self.id().get(store).check_may_leave(caller)?;
 
-        if to_thread.is_none() && !yielding {
-            // This is a `thread.suspend` call
-            store.concurrent_state_mut().check_blocking()?;
+        if to_thread.is_none() {
+            let state = store.concurrent_state_mut();
+            if yielding {
+                // This is a `thread.yield` call
+                if !state.may_block(state.guest_thread.unwrap().task) {
+                    // The spec defines `thread.yield` to be a no-op in a
+                    // non-blocking context, so we return immediately without giving
+                    // any other thread a chance to run.
+                    return Ok(WaitResult::Completed);
+                }
+            } else {
+                // This is a `thread.suspend` call
+                state.check_blocking()?;
+            }
         }
 
         // There could be a pending cancellation from a previous uncancellable wait
@@ -4795,13 +4825,17 @@ impl ConcurrentState {
     }
 
     fn check_blocking_for(&mut self, task: TableId<GuestTask>) -> Result<()> {
-        let task = self.get_mut(task).unwrap();
-        if task.async_function || task.returned_or_cancelled() {
+        if self.may_block(task) {
             Ok(())
         } else {
             Err(Trap::CannotBlockSyncTask.into())
         }
     }
+
+    fn may_block(&mut self, task: TableId<GuestTask>) -> bool {
+        let task = self.get_mut(task).unwrap();
+        task.async_function || task.returned_or_cancelled()
+    }
 }
 
 /// Provide a type hint to compiler about the shape of a parameter lower