trap on blocking call in sync task before return

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 {

crates/test-util/src/wast.rs

@@ -52,6 +52,26 @@ pub fn find_tests(root: &Path) -> Result<Vec<WastTest>> {
         &FindConfig::Infer(component_test_config),
     )
     .with_context(|| format!("failed to add tests from `{}`", cm_tests.display()))?;
+
+    // Temporarily work around upstream tests that loop forever.
+    //
+    // Now that `thread.yield` and `CALLBACK_CODE_YIELD` are both no-ops in
+    // non-blocking contexts, these tests need to be updated; meanwhile, we skip
+    // them.
+    //
+    // TODO: remove this once
+    // https://github.com/WebAssembly/component-model/pull/578 has been merged:
+    {
+        let skip_list = &["drop-subtask.wast", "async-calls-sync.wast"];
+        tests.retain(|test| {
+            test.path
+                .file_name()
+                .and_then(|name| name.to_str())
+                .map(|name| !skip_list.contains(&name))
+                .unwrap_or(true)
+        });
+    }
+
     Ok(tests)
 }
 
@@ -679,6 +699,26 @@ impl WastTest {
             "component-model/test/values/trap-in-post-return.wast",
             "component-model/test/wasmtime/resources.wast",
             "component-model/test/wasm-tools/naming.wast",
+            // TODO: remove these once
+            // https://github.com/WebAssembly/component-model/pull/578 has been
+            // merged:
+            "component-model/test/async/async-calls-sync.wast",
+            "component-model/test/async/backpressure-deadlock.wast",
+            "component-model/test/async/cancel-stream.wast",
+            "component-model/test/async/cancel-subtask.wast",
+            "component-model/test/async/deadlock.wast",
+            "component-model/test/async/drop-subtask.wast",
+            "component-model/test/async/drop-waitable-set.wast",
+            "component-model/test/async/empty-wait.wast",
+            "component-model/test/async/fused.wast",
+            "component-model/test/async/future-read.wast",
+            "component-model/test/async/partial-stream-copies.wast",
+            "component-model/test/async/passing-resources.wast",
+            "component-model/test/async/stackful.wast",
+            "component-model/test/async/trap-if-block-and-sync.wast",
+            "component-model/test/async/trap-if-done.wast",
+            "component-model/test/async/wait-during-callback.wast",
+            "component-model/test/async/zero-length.wast",
         ];
         if failing_component_model_tests
             .iter()