address review feedback

- Assert that `StoreOpaque::suspend` is not called in a non-blocking context except in specific circumstances

- Typecheck async-ness for dynamic host functions

- Use type parameter instead of value parameter in `call_host[_dynamic]`

Signed-off-by: Joel Dice <[email protected]>

Author: dicej

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

@@ -574,6 +574,7 @@ enum SuspendReason {
     Waiting {
         set: TableId<WaitableSet>,
         thread: QualifiedThreadId,
+        skip_may_block_check: bool,
     },
     /// The fiber has finished handling its most recent work item and is waiting
     /// for another (or to be dropped if it is no longer needed).
@@ -582,7 +583,10 @@ enum SuspendReason {
     /// chance to run.
     Yielding { thread: QualifiedThreadId },
     /// The fiber was explicitly suspended with a call to `thread.suspend` or `thread.switch-to`.
-    ExplicitlySuspending { thread: QualifiedThreadId },
+    ExplicitlySuspending {
+        thread: QualifiedThreadId,
+        skip_may_block_check: bool,
+    },
 }
 
 /// Represents a pending call into guest code for a given guest task.
@@ -805,6 +809,7 @@ pub(crate) fn poll_and_block<R: Send + Sync + 'static>(
             store.suspend(SuspendReason::Waiting {
                 set,
                 thread: caller,
+                skip_may_block_check: false,
             })?;
         }
     }
@@ -1445,7 +1450,7 @@ impl StoreOpaque {
                 SuspendReason::ExplicitlySuspending { thread, .. } => {
                     state.get_mut(thread.thread)?.state = GuestThreadState::Suspended(fiber);
                 }
-                SuspendReason::Waiting { set, thread } => {
+                SuspendReason::Waiting { set, thread, .. } => {
                     let old = state
                         .get_mut(set)?
                         .waiting
@@ -1485,6 +1490,26 @@ impl StoreOpaque {
             None
         };
 
+        // We should not have reached here unless either there's no current
+        // task, or the current task is permitted to block.  In addition, we
+        // special-case `thread.switch-to` and waiting for a subtask to go from
+        // `starting` to `started`, both of which we consider non-blocking
+        // operations despite requiring a suspend.
+        assert!(
+            matches!(
+                reason,
+                SuspendReason::ExplicitlySuspending {
+                    skip_may_block_check: true,
+                    ..
+                } | SuspendReason::Waiting {
+                    skip_may_block_check: true,
+                    ..
+                }
+            ) || old_guest_thread
+                .map(|thread| self.concurrent_state_mut().may_block(thread.task))
+                .unwrap_or(true)
+        );
+
         let suspend_reason = &mut self.concurrent_state_mut().suspend_reason;
         assert!(suspend_reason.is_none());
         *suspend_reason = Some(reason);
@@ -1540,6 +1565,7 @@ impl StoreOpaque {
         self.suspend(SuspendReason::Waiting {
             set,
             thread: caller,
+            skip_may_block_check: false,
         })?;
         let state = self.concurrent_state_mut();
         waitable.join(state, old_set)
@@ -2308,6 +2334,7 @@ impl Instance {
         let async_caller = storage.is_none();
         let state = store.0.concurrent_state_mut();
         let guest_thread = state.guest_thread.unwrap();
+        let callee_async = state.get_mut(guest_thread.task)?.async_function;
         let may_enter_after_call = state
             .get_mut(guest_thread.task)?
             .call_post_return_automatically();
@@ -2401,6 +2428,14 @@ impl Instance {
             store.0.suspend(SuspendReason::Waiting {
                 set,
                 thread: caller,
+                // Normally, `StoreOpaque::suspend` would assert it's being
+                // called from a context where blocking is allowed.  However, if
+                // `async_caller` is `true`, we'll only "block" long enough for
+                // the callee to start, i.e. we won't repeat this loop, so we
+                // tell `suspend` it's okay even if we're not allowed to block.
+                // Alternatively, if the callee is not an async function, then
+                // we know it won't block anyway.
+                skip_may_block_check: async_caller || !callee_async,
             })?;
 
             let state = store.0.concurrent_state_mut();
@@ -2885,6 +2920,9 @@ impl Instance {
         self.id().get(store).check_may_leave(caller)?;
 
         if !self.options(store, options).async_ {
+            // The caller may only call `waitable-set.wait` from an async task
+            // (i.e. a task created via a call to an async export).
+            // Otherwise, we'll trap.
             store.concurrent_state_mut().check_blocking()?;
         }
 
@@ -3136,6 +3174,10 @@ impl Instance {
         } else {
             SuspendReason::ExplicitlySuspending {
                 thread: guest_thread,
+                // Tell `StoreOpaque::suspend` it's okay to suspend here since
+                // we're handling a `thread.switch-to` call; otherwise it would
+                // panic if we called it in a non-blocking context.
+                skip_may_block_check: to_thread.is_some(),
             }
         };
 
@@ -3193,6 +3235,7 @@ impl Instance {
                 store.suspend(SuspendReason::Waiting {
                     set,
                     thread: guest_thread,
+                    skip_may_block_check: false,
                 })?;
             }
         }

crates/wasmtime/src/runtime/component/func/host.rs

@@ -126,13 +126,12 @@ impl HostFunc {
         let data = SendSyncPtr::new(NonNull::new(data.as_ptr() as *mut F).unwrap());
         unsafe {
             call_host_and_handle_result::<T>(cx, |store, instance| {
-                call_host(
+                call_host::<_, _, _, _, S>(
                     store,
                     instance,
                     TypeFuncIndex::from_u32(ty),
                     OptionsIndex::from_u32(options),
                     NonNull::slice_from_raw_parts(storage, storage_len).as_mut(),
-                    S::ASYNC,
                     move |store, args| (*data.as_ptr())(store, args),
                 )
             })
@@ -155,12 +154,17 @@ impl HostFunc {
     {
         Arc::new(HostFunc {
             entrypoint: dynamic_entrypoint::<T, F, S>,
-            // This function performs dynamic type checks and subsequently does
-            // not need to perform up-front type checks. Instead everything is
-            // dynamically managed at runtime.
-            //
-            // TODO: Where does async checking happen?
-            typecheck: Box::new(move |_expected_index, _expected_types| Ok(())),
+            // This function performs dynamic type checks on its parameters and
+            // results and subsequently does not need to perform up-front type
+            // checks. However, we _do_ verify async-ness here.
+            typecheck: Box::new(move |ty, types| {
+                let ty = &types.types[ty];
+                if S::ASYNC != ty.async_ {
+                    bail!("type mismatch with async");
+                }
+
+                Ok(())
+            }),
             func: Box::new(func),
         })
     }
@@ -253,26 +257,27 @@ where
 /// * `Return` - the result of the host function
 /// * `F` - the `closure` to actually receive the `Params` and return the
 ///   `Return`
+/// * `S` - the expected `FunctionStyle`
 ///
 /// It's expected that `F` will "un-tuple" the arguments to pass to a host
 /// closure.
 ///
 /// This function is in general `unsafe` as the validity of all the parameters
 /// must be upheld. Generally that's done by ensuring this is only called from
 /// the select few places it's intended to be called from.
-unsafe fn call_host<T, Params, Return, F>(
+unsafe fn call_host<T, Params, Return, F, S>(
     store: StoreContextMut<'_, T>,
     instance: Instance,
     ty: TypeFuncIndex,
     options: OptionsIndex,
     storage: &mut [MaybeUninit<ValRaw>],
-    async_function: bool,
     closure: F,
 ) -> Result<()>
 where
     F: Fn(StoreContextMut<'_, T>, Params) -> HostResult<Return> + Send + Sync + 'static,
     Params: Lift,
     Return: Lower + 'static,
+    S: FunctionStyle,
 {
     let (component, store) = instance.component_and_store_mut(store.0);
     let mut store = StoreContextMut(store);
@@ -369,7 +374,10 @@ where
             );
         }
     } else {
-        if async_function {
+        if S::ASYNC {
+            // The caller has synchronously lowered an async function, meaning
+            // the caller can only call it from an async task (i.e. a task
+            // created via a call to an async export).  Otherwise, we'll trap.
             concurrent::check_blocking(store.0)?;
         }
 
@@ -731,13 +739,12 @@ where
     }
 }
 
-unsafe fn call_host_dynamic<T, F>(
+unsafe fn call_host_dynamic<T, F, S>(
     store: StoreContextMut<'_, T>,
     instance: Instance,
     ty: TypeFuncIndex,
     options: OptionsIndex,
     storage: &mut [MaybeUninit<ValRaw>],
-    async_function: bool,
     closure: F,
 ) -> Result<()>
 where
@@ -751,6 +758,7 @@ where
         + Sync
         + 'static,
     T: 'static,
+    S: FunctionStyle,
 {
     let (component, store) = instance.component_and_store_mut(store.0);
     let mut store = StoreContextMut(store);
@@ -853,7 +861,10 @@ where
             );
         }
     } else {
-        if async_function {
+        if S::ASYNC {
+            // The caller has synchronously lowered an async function, meaning
+            // the caller can only call it from an async task (i.e. a task
+            // created via a call to an async export).  Otherwise, we'll trap.
             concurrent::check_blocking(store.0)?;
         }
 
@@ -975,13 +986,12 @@ where
     let data = SendSyncPtr::new(NonNull::new(data.as_ptr() as *mut F).unwrap());
     unsafe {
         call_host_and_handle_result(cx, |store, instance| {
-            call_host_dynamic::<T, _>(
+            call_host_dynamic::<T, _, S>(
                 store,
                 instance,
                 TypeFuncIndex::from_u32(ty),
                 OptionsIndex::from_u32(options),
                 NonNull::slice_from_raw_parts(storage, storage_len).as_mut(),
-                S::ASYNC,
                 &*data.as_ptr(),
             )
         })