trap on blocking call in sync task before return

This implements a spec change (PR pending) such that tasks created for calls to
synchronous exports may not call potentially-blocking imports or return `wait`
or `poll` callback codes prior to returning a value.  Specifically, the
following are prohibited in that scenario:

- returning callback-code.{wait,poll}
- sync calling async import
- sync calling subtask.cancel
- sync calling {stream,future}.{read,write}
- sync calling {stream,future}.cancel-{read,write}
- calling waitable-set.{wait,poll}
- calling thread.suspend

This breaks a number of tests, which will be addressed in follow-up commits:

- The `{tcp,udp}-socket.bind` implementation in `wasmtime-wasi` is implemented
  using `Linker::func_wrap_concurrent` and thus assumed to be async, whereas the
  WIT interface says they're sync, leading to a type mismatch error at runtime.
  Alex and I have discussed this and have a general plan to address it.

- A number of tests in the tests/component-model submodule that points to the
  spec repo are failing.  Those will presumably be fixed as part of the upcoming
  spec PR (although some could be due to bugs in this implementation, in which
  case I'll fix them).

- A number of tests in tests/misc_testsuite are failing.  I'll address those in
  a follow-up commit.

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

Author: dicej

crates/environ/src/component.rs

@@ -132,6 +132,7 @@ macro_rules! foreach_builtin_component_function {
                 caller_instance: u32,
                 callee_instance: u32,
                 task_return_type: u32,
+                callee_async: u32,
                 string_encoding: u32,
                 result_count_or_max_if_async: u32,
                 storage: ptr_u8,

crates/environ/src/fact.rs

@@ -47,6 +47,7 @@ pub static PREPARE_CALL_FIXED_PARAMS: &[ValType] = &[
     ValType::I32,     // caller_instance
     ValType::I32,     // callee_instance
     ValType::I32,     // task_return_type
+    ValType::I32,     // callee_async
     ValType::I32,     // string_encoding
     ValType::I32,     // result_count_or_max_if_async
 ];

crates/environ/src/fact/trampoline.rs

@@ -544,6 +544,11 @@ impl<'a, 'b> Compiler<'a, 'b> {
         self.instruction(I32Const(
             i32::try_from(self.types[adapter.lift.ty].results.as_u32()).unwrap(),
         ));
+        self.instruction(I32Const(if self.types[adapter.lift.ty].async_ {
+            1
+        } else {
+            0
+        }));
         self.instruction(I32Const(i32::from(
             adapter.lift.options.string_encoding as u8,
         )));

crates/environ/src/trap_encoding.rs

@@ -112,6 +112,9 @@ pub enum Trap {
     /// scenario where a component instance tried to call an import or intrinsic
     /// when it wasn't allowed to, e.g. from a post-return function.
     CannotLeaveComponent,
+
+    /// A synchronous task attempted to make a potentially blocking call.
+    CannotBlockSyncTask,
     // if adding a variant here be sure to update the `check!` macro below
 }
 
@@ -154,6 +157,7 @@ impl Trap {
             DisabledOpcode
             AsyncDeadlock
             CannotLeaveComponent
+            CannotBlockSyncTask
         }
 
         None
@@ -190,6 +194,7 @@ impl fmt::Display for Trap {
             DisabledOpcode => "pulley opcode disabled at compile time was executed",
             AsyncDeadlock => "deadlock detected: event loop cannot make further progress",
             CannotLeaveComponent => "cannot leave component instance",
+            CannotBlockSyncTask => "cannot block a synchronous task",
         };
         write!(f, "wasm trap: {desc}")
     }

crates/misc/component-async-tests/wit/test.wit

@@ -122,7 +122,7 @@ interface resource-stream {
     foo: func();
   }
 
-  foo: func(count: u32) -> stream<x>;
+  foo: async func(count: u32) -> stream<x>;
 }
 
 interface closed {
@@ -157,7 +157,7 @@ interface cancel {
     leak-task-after-cancel,
   }
 
-  run: func(mode: mode, cancel-delay-millis: u64);
+  run: async func(mode: mode, cancel-delay-millis: u64);
 }
 
 interface intertask {

crates/test-programs/src/bin/async_read_resource_stream.rs

@@ -15,7 +15,7 @@ struct Component;
 impl Guest for Component {
     async fn run() {
         let mut count = 7;
-        let mut stream = resource_stream::foo(count);
+        let mut stream = resource_stream::foo(count).await;
 
         while let Some(x) = stream.next().await {
             if count > 0 {

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

@@ -705,6 +705,12 @@ pub(crate) enum WaitResult {
     Completed,
 }
 
+/// Raise a trap if the calling task is synchronous and trying to block prior to
+/// returning a value.
+pub(crate) fn check_blocking(store: &mut dyn VMStore) -> Result<()> {
+    store.concurrent_state_mut().check_blocking()
+}
+
 /// Poll the specified future until it completes on behalf of a guest->host call
 /// using a sync-lowered import.
 ///
@@ -1643,6 +1649,8 @@ impl Instance {
                 }));
             }
             callback_code::WAIT | callback_code::POLL => {
+                state.check_blocking_for(guest_thread.task)?;
+
                 let set = get_set(store, set)?;
                 let state = store.concurrent_state_mut();
 
@@ -2038,6 +2046,7 @@ impl Instance {
         caller_instance: RuntimeComponentInstanceIndex,
         callee_instance: RuntimeComponentInstanceIndex,
         task_return_type: TypeTupleIndex,
+        callee_async: bool,
         memory: *mut VMMemoryDefinition,
         string_encoding: u8,
         caller_info: CallerInfo,
@@ -2181,6 +2190,7 @@ impl Instance {
             },
             None,
             callee_instance,
+            callee_async,
         )?;
 
         let guest_task = state.push(new_task)?;
@@ -2848,6 +2858,10 @@ impl Instance {
         set: u32,
         payload: u32,
     ) -> Result<u32> {
+        if !self.options(store, options).async_ {
+            store.concurrent_state_mut().check_blocking()?;
+        }
+
         self.id().get(store).check_may_leave(caller)?;
         let &CanonicalOptions {
             cancellable,
@@ -2877,6 +2891,10 @@ impl Instance {
         set: u32,
         payload: u32,
     ) -> Result<u32> {
+        if !self.options(store, options).async_ {
+            store.concurrent_state_mut().check_blocking()?;
+        }
+
         self.id().get(store).check_may_leave(caller)?;
         let &CanonicalOptions {
             cancellable,
@@ -3056,6 +3074,11 @@ impl Instance {
         yielding: bool,
         to_thread: Option<u32>,
     ) -> Result<WaitResult> {
+        if to_thread.is_none() && !yielding {
+            // This is a `thread.suspend` call
+            store.concurrent_state_mut().check_blocking()?;
+        }
+
         // There could be a pending cancellation from a previous uncancellable wait
         if cancellable && store.concurrent_state_mut().take_pending_cancellation() {
             return Ok(WaitResult::Cancelled);
@@ -3185,6 +3208,10 @@ impl Instance {
         async_: bool,
         task_id: u32,
     ) -> Result<u32> {
+        if !async_ {
+            store.concurrent_state_mut().check_blocking()?;
+        }
+
         self.id().get(store).check_may_leave(caller_instance)?;
         let (rep, is_host) =
             self.id().get_mut(store).guest_tables().0[caller_instance].subtask_rep(task_id)?;
@@ -3345,6 +3372,7 @@ pub trait VMComponentAsyncStore {
         caller_instance: RuntimeComponentInstanceIndex,
         callee_instance: RuntimeComponentInstanceIndex,
         task_return_type: TypeTupleIndex,
+        callee_async: bool,
         string_encoding: u8,
         result_count: u32,
         storage: *mut ValRaw,
@@ -3504,6 +3532,7 @@ impl<T: 'static> VMComponentAsyncStore for StoreInner<T> {
         caller_instance: RuntimeComponentInstanceIndex,
         callee_instance: RuntimeComponentInstanceIndex,
         task_return_type: TypeTupleIndex,
+        callee_async: bool,
         string_encoding: u8,
         result_count_or_max_if_async: u32,
         storage: *mut ValRaw,
@@ -3522,6 +3551,7 @@ impl<T: 'static> VMComponentAsyncStore for StoreInner<T> {
                 caller_instance,
                 callee_instance,
                 task_return_type,
+                callee_async,
                 memory,
                 string_encoding,
                 match result_count_or_max_if_async {
@@ -4062,6 +4092,9 @@ pub(crate) struct GuestTask {
     /// The state of the host future that represents an async task, which must
     /// be dropped before we can delete the task.
     host_future_state: HostFutureState,
+    /// Indicates whether this task was created for a call to an async-lifted
+    /// export.
+    async_function: bool,
 }
 
 impl GuestTask {
@@ -4102,6 +4135,7 @@ impl GuestTask {
         caller: Caller,
         callback: Option<CallbackFn>,
         component_instance: RuntimeComponentInstanceIndex,
+        async_function: bool,
     ) -> Result<Self> {
         let sync_call_set = state.push(WaitableSet::default())?;
         let host_future_state = match &caller {
@@ -4136,6 +4170,7 @@ impl GuestTask {
             exited: false,
             threads: HashSet::new(),
             host_future_state,
+            async_function,
         })
     }
 
@@ -4749,6 +4784,20 @@ impl ConcurrentState {
             false
         }
     }
+
+    fn check_blocking(&mut self) -> Result<()> {
+        let task = self.guest_thread.unwrap().task;
+        self.check_blocking_for(task)
+    }
+
+    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()) {
+            Err(Trap::CannotBlockSyncTask.into())
+        } else {
+            Ok(())
+        }
+    }
 }
 
 /// Provide a type hint to compiler about the shape of a parameter lower
@@ -4907,7 +4956,9 @@ pub(crate) fn prepare_call<T, R>(
 
     let instance = handle.instance().id().get(store.0);
     let options = &instance.component().env_component().options[options];
-    let task_return_type = instance.component().types()[ty].results;
+    let ty = &instance.component().types()[ty];
+    let async_function = ty.async_;
+    let task_return_type = ty.results;
     let component_instance = raw_options.instance;
     let callback = options.callback.map(|i| instance.runtime_callback(i));
     let memory = options
@@ -4964,6 +5015,7 @@ pub(crate) fn prepare_call<T, R>(
             ) as CallbackFn
         }),
         component_instance,
+        async_function,
     )?;
     task.function_index = Some(handle.index());
 

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

@@ -3093,6 +3093,10 @@ impl Instance {
         address: u32,
         count: u32,
     ) -> Result<ReturnCode> {
+        if !self.options(store.0, options).async_ {
+            store.0.concurrent_state_mut().check_blocking()?;
+        }
+
         let address = usize::try_from(address).unwrap();
         let count = usize::try_from(count).unwrap();
         self.check_bounds(store.0, options, ty, address, count)?;
@@ -3315,6 +3319,10 @@ impl Instance {
         address: u32,
         count: u32,
     ) -> Result<ReturnCode> {
+        if !self.options(store.0, options).async_ {
+            store.0.concurrent_state_mut().check_blocking()?;
+        }
+
         let address = usize::try_from(address).unwrap();
         let count = usize::try_from(count).unwrap();
         self.check_bounds(store.0, options, ty, address, count)?;
@@ -3686,6 +3694,10 @@ impl Instance {
         async_: bool,
         writer: u32,
     ) -> Result<ReturnCode> {
+        if !async_ {
+            store.concurrent_state_mut().check_blocking()?;
+        }
+
         let (rep, state) =
             get_mut_by_index_from(self.id().get_mut(store).table_for_transmit(ty), ty, writer)?;
         let id = TableId::<TransmitHandle>::new(rep);
@@ -3720,6 +3732,10 @@ impl Instance {
         async_: bool,
         reader: u32,
     ) -> Result<ReturnCode> {
+        if !async_ {
+            store.concurrent_state_mut().check_blocking()?;
+        }
+
         let (rep, state) =
             get_mut_by_index_from(self.id().get_mut(store).table_for_transmit(ty), ty, reader)?;
         let id = TableId::<TransmitHandle>::new(rep);