fix(bindgen): return JSPI host imports synchronously when the context cannot wait

crates/js-component-bindgen/src/core.rs

@@ -55,6 +55,62 @@ use wasmtime_environ::{
     PrimaryMap,
 };
 
+/// Dispatch a host import through JSPI only when its caller may block. The
+/// dispatcher must itself be Wasm: a JS wrapper would prevent suspension across
+/// the intervening JS frame. Preserve the exact canonical lowering signature.
+pub(crate) fn conditional_host_import(ty: &wasmtime_environ::WasmFuncType) -> Vec<u8> {
+    use wasm_encoder::{BlockType, Instruction, ValType};
+    use wasmtime_environ::WasmValType;
+
+    let encode_type = |ty: &WasmValType| match ty {
+        WasmValType::I32 => ValType::I32,
+        WasmValType::I64 => ValType::I64,
+        WasmValType::F32 => ValType::F32,
+        WasmValType::F64 => ValType::F64,
+        _ => unreachable!("host import canonical ABI must use numeric types"),
+    };
+    let mut types = TypeSection::new();
+    types.ty().function([], [ValType::I32]);
+    types.ty().function(
+        ty.params().iter().map(encode_type),
+        ty.results().iter().map(encode_type),
+    );
+
+    let mut imports = ImportSection::new();
+    imports.import("", "mayBlock", EntityType::Function(0));
+    imports.import("", "fast", EntityType::Function(1));
+    imports.import("", "slow", EntityType::Function(1));
+    let mut functions = FunctionSection::new();
+    functions.function(1);
+    let mut exports = ExportSection::new();
+    exports.export("run", ExportKind::Func, 3);
+
+    let mut body = Function::new([]);
+    body.instruction(&Instruction::Call(0));
+    body.instruction(&Instruction::If(BlockType::Empty));
+    for i in 0..ty.params().len() {
+        body.instruction(&Instruction::LocalGet(i as u32));
+    }
+    body.instruction(&Instruction::Call(2));
+    body.instruction(&Instruction::Return);
+    body.instruction(&Instruction::End);
+    for i in 0..ty.params().len() {
+        body.instruction(&Instruction::LocalGet(i as u32));
+    }
+    body.instruction(&Instruction::Call(1));
+    body.instruction(&Instruction::End);
+    let mut code = CodeSection::new();
+    code.function(&body);
+
+    let mut module = Module::new();
+    module.section(&types);
+    module.section(&imports);
+    module.section(&functions);
+    module.section(&exports);
+    module.section(&code);
+    module.finish()
+}
+
 pub enum Translation<'a> {
     Normal(ModuleTranslation<'a>),
     Augmented {

crates/js-component-bindgen/src/function_bindgen.rs

@@ -229,6 +229,11 @@ pub struct FunctionBindgen<'a> {
     /// Whether the function is guest async lifted (i.e. WASI P3)
     pub is_async: bool,
 
+    /// Whether the body of an import is generated for a generator function (run by
+    /// `MaybeSyncGeneratorFn`), which yields the host's answer instead of awaiting it,
+    /// so that it is synchronous unless the host returns a promise
+    pub(crate) host_import_generator: bool,
+
     /// Whether the canonical ABI call itself is async.
     ///
     /// This can differ from `is_async` when an async component function is
@@ -2391,12 +2396,19 @@ impl Bindgen for FunctionBindgen<'_> {
                 // at the CM level -- async functions always return
                 let fn_wasm_result_count = if func.result.is_none() { 0 } else { 1 };
 
+                // A generator yields only the host's answer, and a host task enters at once
+                let (await_, enter_task) = if self.host_import_generator {
+                    ("yield ", "task.enterHost()")
+                } else {
+                    ("await ", "await task.enter({ isHost: hostProvided })")
+                };
+
                 // If the task is async, do an explicit wait for backpressure before the call execution
                 if task_managed && is_async {
                     uwriteln!(
                         self.src,
                         r#"
-                        const started = await task.enter({{ isHost: hostProvided }});
+                        const started = {enter_task};
                         if (!started) {{
                             {debug_log_fn}('[Instruction::CallInterface] failed to enter task', {{
                                 taskID: task.id(),
@@ -2424,7 +2436,7 @@ impl Bindgen for FunctionBindgen<'_> {
                     || self.requires_async_porcelain
                 {
                     (
-                        "await ",
+                        await_,
                         self.task_intrinsic(Intrinsic::WithGlobalCurrentThreadMetaFnAsync),
                         format!(
                             r#"

crates/js-component-bindgen/src/intrinsics/mod.rs

@@ -169,9 +169,16 @@ pub enum Intrinsic {
     /// Execute a closure with a certain set current thread
     WithGlobalCurrentThreadMetaFn,
 
-    /// Execute an async closure with a certain set current thread
+    /// Execute a closure with a certain set current thread, which is async only
+    /// if the closure returns a promise
     WithGlobalCurrentThreadMetaFnAsync,
 
+    /// Run a generator function synchronously, until it yields a promise
+    MaybeSyncGeneratorFn,
+
+    /// Bypass the JSPI boundary when a host import's caller cannot block.
+    ConditionalSuspendingHostImportFn,
+
     /// Clear the global thread meta
     ClearGlobalCurrentThreadMetaFn,
 
@@ -1212,6 +1219,8 @@ impl Intrinsic {
             // In the future, this mechanism may be replaced with a simple event loop that necessarily executes
             // all pending work serially, with this intrinsic becoming simply queueing work onto that event loop.
             //
+            // A closure that returns synchronously is treated like WithGlobalCurrentThreadMetaFn.
+            // While the promise of one that returns a promise is pending, the task stays current.
             Self::WithGlobalCurrentThreadMetaFnAsync => {
                 let debug_log_fn = args.require_intrinsic(Intrinsic::DebugLog);
                 let with_global_current_thread_meta_async_fn =
@@ -1221,7 +1230,7 @@ impl Intrinsic {
 
                 output.push_str(&format!(
                     r#"
-                      async function {with_global_current_thread_meta_async_fn}(args) {{
+                      function {with_global_current_thread_meta_async_fn}(args) {{
                           {debug_log_fn}('[{with_global_current_thread_meta_async_fn}()] args', args);
                           if (!args) {{ throw new TypeError('args missing'); }}
                           if (args.taskID === undefined) {{ throw new TypeError('missing task ID'); }}
@@ -1230,27 +1239,88 @@ impl Intrinsic {
                           if (!args.fn) {{ throw new TypeError('missing fn'); }}
 
                           const {{ taskID, threadID, componentIdx, fn }} = args;
+                          const previous = {global_current_thread_meta_obj}[componentIdx] ?? null;
+                          const previousCurrent = {global_current_thread_meta_obj}.current ?? null;
 
+                          let result;
                           try {{
                               {global_current_thread_meta_obj}.current =
-                                  {global_current_thread_meta_obj}[componentIdx] = {{ threadID, taskID, componentIdx }};
-                              return await fn();
+                                  {global_current_thread_meta_obj}[componentIdx] = {{ taskID, threadID, componentIdx }};
+                              result = fn();
                           }} catch (err) {{
-                              {debug_log_fn}("error while executing async callee/callback", {{
+                              {debug_log_fn}("error while executing sync callee/callback", {{
                                   ...args,
                                   err,
                               }});
+                              {global_current_thread_meta_obj}[componentIdx] = previous;
+                              {global_current_thread_meta_obj}.current = previousCurrent;
                               throw err;
-                          }} finally {{
-                              if ({global_current_thread_meta_obj}[componentIdx]?.threadID === threadID &&
-                                  {global_current_thread_meta_obj}[componentIdx]?.taskID === taskID) {{
-                                  {global_current_thread_meta_obj}[componentIdx] = null;
-                              }}
-                              if ({global_current_thread_meta_obj}.current?.threadID === threadID &&
-                                  {global_current_thread_meta_obj}.current?.taskID === taskID) {{
-                                  {global_current_thread_meta_obj}.current = null;
-                              }}
                           }}
+
+                          if (result === null ||
+                              (typeof result !== 'object' && typeof result !== 'function') ||
+                              typeof result.then !== 'function') {{
+                              {global_current_thread_meta_obj}[componentIdx] = previous;
+                              {global_current_thread_meta_obj}.current = previousCurrent;
+                              return result;
+                          }}
+
+                          return (async () => {{
+                              try {{
+                                  return await result;
+                              }} catch (err) {{
+                                  {debug_log_fn}("error while executing async callee/callback", {{
+                                      ...args,
+                                      err,
+                                  }});
+                                  throw err;
+                              }} finally {{
+                                  if ({global_current_thread_meta_obj}[componentIdx]?.threadID === threadID &&
+                                      {global_current_thread_meta_obj}[componentIdx]?.taskID === taskID) {{
+                                      {global_current_thread_meta_obj}[componentIdx] = null;
+                                  }}
+                                  if ({global_current_thread_meta_obj}.current?.threadID === threadID &&
+                                      {global_current_thread_meta_obj}.current?.taskID === taskID) {{
+                                      {global_current_thread_meta_obj}.current = null;
+                                  }}
+                              }}
+                          }})();
+                      }}
+                    "#,
+                ));
+            }
+
+            // Wrap a generator function, which yields where an async function would await,
+            // so that it runs synchronously until it yields a promise. A host import that
+            // answers synchronously then returns synchronously, which a task that cannot
+            // block (e.g. a resource destructor) requires.
+            Self::MaybeSyncGeneratorFn => {
+                let maybe_sync_generator_fn = args.require_intrinsic(Self::MaybeSyncGeneratorFn);
+
+                output.push_str(&format!(
+                    r#"
+                      function {maybe_sync_generator_fn}(genFn) {{
+                          return function (...args) {{
+                              const gen = genFn.apply(this, args);
+                              const step = (method, value) => {{
+                                  while (true) {{
+                                      const {{ value: yielded, done }} = gen[method](value);
+                                      if (done) {{ return yielded; }}
+                                      if (yielded === null ||
+                                          (typeof yielded !== 'object' && typeof yielded !== 'function') ||
+                                          typeof yielded.then !== 'function') {{
+                                          method = 'next';
+                                          value = yielded;
+                                          continue;
+                                      }}
+                                      return Promise.resolve(yielded).then(
+                                          res => step('next', res),
+                                          err => step('throw', err),
+                                      );
+                                  }}
+                              }};
+                              return step('next', undefined);
+                          }};
                       }}
                     "#,
                 ));
@@ -1300,6 +1370,8 @@ impl Intrinsic {
             // can abandon an unobserved promise. A nonblocking task may still
             // call a suspending import when the helper completes synchronously,
             // so preserve its raw result and reject only a returned promise.
+            // A host import (`callsHost`) handles a call without a current task,
+            // such as a destructor run outside of any task, in the same way.
             //
             // Run blocking helpers entirely inside an async function. Moving
             // the helper invocation outside that function changes JSPI
@@ -1327,7 +1399,7 @@ impl Intrinsic {
 
                 output.push_str(&format!(
                     r#"
-                      function {suspending_import_wrapper_fn}(componentIdx, fn, syncOnly = false, switchesTask = false, convertsRangeErrors = true) {{
+                      function {suspending_import_wrapper_fn}(componentIdx, fn, syncOnly = false, switchesTask = false, convertsRangeErrors = true, callsHost = false) {{
                           return function (...args) {{
                               {check_may_leave_fn}(componentIdx);
                               const saved = {global_current_thread_meta_obj}[componentIdx] ?? null;
@@ -1338,7 +1410,7 @@ impl Intrinsic {
                               const mayBlock = savedThread
                                   ? (savedThread?.task.mayBlock() ?? ({current_task_may_block}.value !== 0))
                                   : false;
-                              if (!saved && !mayBlock) {{
+                              if (!saved && !mayBlock && !callsHost) {{
                                   throw new {runtime_error_class}('cannot block a synchronous task before returning');
                               }}
 
@@ -1389,6 +1461,30 @@ impl Intrinsic {
                 ));
             }
 
+            Self::ConditionalSuspendingHostImportFn => {
+                let wrapper = args.require_intrinsic(*self);
+                let current_meta = args.require_intrinsic(Self::GlobalCurrentThreadMeta);
+                let get_thread = args.require_intrinsic(AsyncTaskIntrinsic::GetCurrentThread);
+                uwriteln!(
+                    output,
+                    r#"
+                    function {wrapper}(module, componentIdx, fn) {{
+                        return new WebAssembly.Instance(module, {{ '': {{
+                            mayBlock: () => {{
+                                const saved = {current_meta}[componentIdx];
+                                const thread = saved
+                                    ? {get_thread}(saved.componentIdx, saved.threadID)
+                                    : null;
+                                return thread?.task.mayBlock() ?? false;
+                            }},
+                            fast: fn,
+                            slow: new WebAssembly.Suspending(fn),
+                        }} }}).exports.run;
+                    }}
+                    "#,
+                );
+            }
+
             Self::ConditionalSuspending1I32ToI32Fn
             | Self::ConditionalSuspending2I32ToI32Fn
             | Self::ConditionalSuspending3I32ToI32Fn
@@ -1661,7 +1757,7 @@ mod tests {
         assert!(source.contains("if (syncOnly || !mayBlock) {"));
 
         assert!(!source.contains("return async function (...args) {"));
-        assert!(source.contains("if (!saved && !mayBlock) {"));
+        assert!(source.contains("if (!saved && !mayBlock && !callsHost) {"));
 
         assert!(source.contains("result = fn.apply(null, args);"));
         assert!(source.contains("typeof result.then === 'function'"));
@@ -1675,6 +1771,26 @@ mod tests {
         assert!(source.contains("CURRENT_TASK_META[componentIdx] = saved;"));
     }
 
+    #[test]
+    fn maybe_sync_generator_awaits_only_promises() {
+        let source = render_intrinsic_body(Intrinsic::MaybeSyncGeneratorFn);
+
+        assert!(source.contains("const gen = genFn.apply(this, args);"));
+        assert!(source.contains("typeof yielded.then !== 'function'"));
+        assert!(source.contains("return Promise.resolve(yielded).then("));
+        assert!(source.contains("err => step('throw', err),"));
+    }
+
+    #[test]
+    fn async_current_task_wrapper_restores_sync_results() {
+        let source = render_intrinsic_body(Intrinsic::WithGlobalCurrentThreadMetaFnAsync);
+
+        assert!(source.contains("result = fn();"));
+        assert!(source.contains("CURRENT_TASK_META[componentIdx] = previous;"));
+        assert!(source.contains("return await result;"));
+        assert!(source.contains("CURRENT_TASK_META[componentIdx] = null;"));
+    }
+
     #[test]
     fn manually_async_tasks_may_block() {
         let source =
@@ -2936,6 +3052,8 @@ impl Intrinsic {
             Self::SetGlobalCurrentThreadMetaFn => "_setGlobalCurrentThreadMeta",
             Self::WithGlobalCurrentThreadMetaFn => "_withGlobalCurrentThreadMeta",
             Self::WithGlobalCurrentThreadMetaFnAsync => "_withGlobalCurrentThreadMetaAsync",
+            Self::MaybeSyncGeneratorFn => "_maybeSyncGenerator",
+            Self::ConditionalSuspendingHostImportFn => "_conditionalSuspendingHostImport",
             Self::ClearGlobalCurrentThreadMetaFn => "_clearCurrentTask",
             Self::SuspendingImportWrapperFn => "_suspendingImport",
             Self::ConditionalSuspending1I32ToI32Fn => "_conditionalSuspending1I32ToI32",