trap on blocking call in sync task before return (#12043)

* 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 an 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]>

* call `check_may_leave` before `check_blocking`

`check_blocking` needs access to the current task, but that's not set for
post-return functions since those should not be calling _any_ imports at all, so
first check for that.

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

* fix `misc_testsuite` test regressions

This amounts to adding `async` to any exported component functions that might
need to block.

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

* simplify code in `ConcurrentState::check_blocking`

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

* 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 <[email protected]>

* fix build when `component-model-async` feature disabled

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

* fix more test regressions

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

* fix more test regressions

Note that this temporarily updates the `tests/component-model` submodule to the
branch for WebAssembly/component-model#577 until that PR
is merged.

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

* tweak `Trap::CannotBlockSyncTask` message

This clarifies that such a task cannot block prior to returning.

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

* fix cancel_host_future test

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

* trap sync-lowered, guest->guest async calls in sync tasks

I somehow forgot to address this earlier.  Thanks to Luke for catching this.

Note that this commit doesn't include test coverage, but Luke's forthecoming
tests in the `component-model` repo will cover it, and we'll pull that in with
the next submodule update.

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

* switch back to `main` branch of `component-model` repo

...and skip or `should_fail` the tests that won't pass until
WebAssembly/component-model#578 is merged.

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

* add `trap-if-block-and-sync.wast`

We'll remove this again in favor of the upstream version once
WebAssembly/component-model#578 has been merged.

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

* 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]>

* add explanation comments to `check_blocking` calls

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

* fix fuzz test oracle for async functions

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

---------

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

Author: dicej

crates/cranelift/src/compiler/component.rs

@@ -742,6 +742,14 @@ impl<'a> TrampolineCompiler<'a> {
                     |_, _| {},
                 );
             }
+            Trampoline::CheckBlocking => {
+                self.translate_libcall(
+                    host::check_blocking,
+                    TrapSentinel::Falsy,
+                    WasmArgs::InRegisters,
+                    |_, _| {},
+                );
+            }
             Trampoline::ContextGet { instance, slot } => {
                 self.translate_libcall(
                     host::context_get,

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,
@@ -186,6 +187,8 @@ macro_rules! foreach_builtin_component_function {
             #[cfg(feature = "component-model-async")]
             error_context_transfer(vmctx: vmctx, src_idx: u32, src_table: u32, dst_table: u32) -> u64;
             #[cfg(feature = "component-model-async")]
+            check_blocking(vmctx: vmctx) -> bool;
+            #[cfg(feature = "component-model-async")]
             context_get(vmctx: vmctx, caller_instance: u32, slot: u32) -> u64;
             #[cfg(feature = "component-model-async")]
             context_set(vmctx: vmctx, caller_instance: u32, slot: u32, val: u32) -> bool;

crates/environ/src/component/dfg.rs

@@ -478,6 +478,7 @@ pub enum Trampoline {
     FutureTransfer,
     StreamTransfer,
     ErrorContextTransfer,
+    CheckBlocking,
     ContextGet {
         instance: RuntimeComponentInstanceIndex,
         slot: u32,
@@ -1160,6 +1161,7 @@ impl LinearizeDfg<'_> {
             Trampoline::FutureTransfer => info::Trampoline::FutureTransfer,
             Trampoline::StreamTransfer => info::Trampoline::StreamTransfer,
             Trampoline::ErrorContextTransfer => info::Trampoline::ErrorContextTransfer,
+            Trampoline::CheckBlocking => info::Trampoline::CheckBlocking,
             Trampoline::ContextGet { instance, slot } => info::Trampoline::ContextGet {
                 instance: *instance,
                 slot: *slot,

crates/environ/src/component/info.rs

@@ -1112,6 +1112,10 @@ pub enum Trampoline {
     /// component does not invalidate the handle in the original component.
     ErrorContextTransfer,
 
+    /// An intrinsic used by FACT-generated modules to check whether an
+    /// async-typed function may be called via a sync lower.
+    CheckBlocking,
+
     /// Intrinsic used to implement the `context.get` component model builtin.
     ///
     /// The payload here represents that this is accessing the Nth slot of local
@@ -1242,6 +1246,7 @@ impl Trampoline {
             FutureTransfer => format!("future-transfer"),
             StreamTransfer => format!("stream-transfer"),
             ErrorContextTransfer => format!("error-context-transfer"),
+            CheckBlocking => format!("check-blocking"),
             ContextGet { .. } => format!("context-get"),
             ContextSet { .. } => format!("context-set"),
             ThreadIndex => format!("thread-index"),

crates/environ/src/component/translate/adapt.rs

@@ -345,6 +345,7 @@ fn fact_import_to_core_def(
         fact::Import::ErrorContextTransfer => {
             simple_intrinsic(dfg::Trampoline::ErrorContextTransfer)
         }
+        fact::Import::CheckBlocking => simple_intrinsic(dfg::Trampoline::CheckBlocking),
     }
 }
 

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
 ];
@@ -89,6 +90,8 @@ pub struct Module<'a> {
     imported_stream_transfer: Option<FuncIndex>,
     imported_error_context_transfer: Option<FuncIndex>,
 
+    imported_check_blocking: Option<FuncIndex>,
+
     // Current status of index spaces from the imports generated so far.
     imported_funcs: PrimaryMap<FuncIndex, Option<CoreDef>>,
     imported_memories: PrimaryMap<MemoryIndex, CoreDef>,
@@ -259,6 +262,7 @@ impl<'a> Module<'a> {
             imported_future_transfer: None,
             imported_stream_transfer: None,
             imported_error_context_transfer: None,
+            imported_check_blocking: None,
             exports: Vec::new(),
         }
     }
@@ -712,6 +716,17 @@ impl<'a> Module<'a> {
         )
     }
 
+    fn import_check_blocking(&mut self) -> FuncIndex {
+        self.import_simple(
+            "async",
+            "check-blocking",
+            &[],
+            &[],
+            Import::CheckBlocking,
+            |me| &mut me.imported_check_blocking,
+        )
+    }
+
     fn translate_helper(&mut self, helper: Helper) -> FunctionId {
         *self.helper_funcs.entry(helper).or_insert_with(|| {
             // Generate a fresh `Function` with a unique id for what we're about to
@@ -870,6 +885,9 @@ pub enum Import {
     /// An intrinisic used by FACT-generated modules to (partially or entirely) transfer
     /// ownership of an `error-context`.
     ErrorContextTransfer,
+    /// An intrinsic used by FACT-generated modules to check whether an
+    /// async-typed function may be called via a sync lower.
+    CheckBlocking,
 }
 
 impl Options {

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,
         )));
@@ -748,6 +753,11 @@ impl<'a, 'b> Compiler<'a, 'b> {
             );
         }
 
+        if self.types[adapter.lift.ty].async_ {
+            let check_blocking = self.module.import_check_blocking();
+            self.instruction(Call(check_blocking.as_u32()));
+        }
+
         if self.emit_resource_call {
             let enter = self.module.import_resource_enter_call();
             self.instruction(Call(enter.as_u32()));

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 before returning",
         };
         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 {