allow instance to be reentered while yielding

crates/cranelift/src/compiler/component.rs

@@ -1,6 +1,6 @@
 //! Compilation support for the component model.
 
-use crate::{TRAP_ALWAYS, TRAP_CANNOT_ENTER, TRAP_INTERNAL_ASSERT, compiler::Compiler};
+use crate::{TRAP_ALWAYS, TRAP_INTERNAL_ASSERT, compiler::Compiler};
 use cranelift_codegen::ir::condcodes::IntCC;
 use cranelift_codegen::ir::{self, InstBuilder, MemFlags, Value};
 use cranelift_codegen::isa::{CallConv, TargetIsa};
@@ -732,9 +732,17 @@ impl<'a> TrampolineCompiler<'a> {
                     |_, _| {},
                 );
             }
-            Trampoline::CheckBlocking => {
+            Trampoline::EnterSyncCall => {
                 self.translate_libcall(
-                    host::check_blocking,
+                    host::enter_sync_call,
+                    TrapSentinel::Falsy,
+                    WasmArgs::InRegisters,
+                    |_, _| {},
+                );
+            }
+            Trampoline::ExitSyncCall => {
+                self.translate_libcall(
+                    host::exit_sync_call,
                     TrapSentinel::Falsy,
                     WasmArgs::InRegisters,
                     |_, _| {},
@@ -1138,11 +1146,9 @@ impl<'a> TrampolineCompiler<'a> {
         //      brif should_run_destructor, run_destructor_block, return_block
         //
         //    run_destructor_block:
-        //      ;; test may_enter, but only if the component instances
+        //      ;; call `enter_sync_call`, but only if the component instances
         //      ;; differ
-        //      flags = load.i32 vmctx+$offset
-        //      masked = band flags, $FLAG_MAY_ENTER
-        //      trapz masked, CANNOT_ENTER_CODE
+        //      ...
         //
         //      ;; ============================================================
         //      ;; this is conditionally emitted based on whether the resource
@@ -1156,6 +1162,9 @@ impl<'a> TrampolineCompiler<'a> {
         //      call_indirect func_addr, callee_vmctx, vmctx, rep
         //      ;; ============================================================
         //
+        //      ;; call `exit_sync_call` if the component instances differ
+        //      ...
+        //
         //      jump return_block
         //
         //    return_block:
@@ -1186,25 +1195,32 @@ impl<'a> TrampolineCompiler<'a> {
         self.builder.switch_to_block(run_destructor_block);
 
         // If this is a defined resource within the component itself then a
-        // check needs to be emitted for the `may_enter` flag. Note though
+        // check needs to be emitted for recursive reentrance. Note though
         // that this check can be elided if the resource table resides in
         // the same component instance that defined the resource as the
         // component is calling itself.
-        if let Some(def) = resource_def {
+        let emit_exit = if let Some(def) = resource_def {
             if self.types[resource].unwrap_concrete_instance() != def.instance {
-                let flags = self.builder.ins().load(
-                    ir::types::I32,
-                    trusted,
+                let host_args = vec![
                     vmctx,
-                    i32::try_from(self.offsets.instance_flags(def.instance)).unwrap(),
-                );
-                let masked = self
-                    .builder
-                    .ins()
-                    .band_imm(flags, i64::from(FLAG_MAY_ENTER));
-                self.builder.ins().trapz(masked, TRAP_CANNOT_ENTER);
+                    self.builder
+                        .ins()
+                        .iconst(ir::types::I32, i64::from(instance.as_u32())),
+                    self.builder.ins().iconst(ir::types::I32, i64::from(0)),
+                    self.builder
+                        .ins()
+                        .iconst(ir::types::I32, i64::from(def.instance.as_u32())),
+                ];
+                let call = self.call_libcall(vmctx, host::enter_sync_call, &host_args);
+                let result = self.builder.func.dfg.inst_results(call).get(0).copied();
+                self.raise_if_host_trapped(result.unwrap());
+                true
+            } else {
+                false
             }
-        }
+        } else {
+            false
+        };
 
         // Conditionally emit destructor-execution code based on whether we
         // statically know that a destructor exists or not.
@@ -1253,6 +1269,13 @@ impl<'a> TrampolineCompiler<'a> {
                 &[callee_vmctx, caller_vmctx, rep],
             );
         }
+
+        if emit_exit {
+            let call = self.call_libcall(vmctx, host::exit_sync_call, &[vmctx]);
+            let result = self.builder.func.dfg.inst_results(call).get(0).copied();
+            self.raise_if_host_trapped(result.unwrap());
+        }
+
         self.builder.ins().jump(return_block, &[]);
         self.builder.seal_block(run_destructor_block);
 

crates/environ/src/component.rs

@@ -99,6 +99,9 @@ macro_rules! foreach_builtin_component_function {
             resource_enter_call(vmctx: vmctx);
             resource_exit_call(vmctx: vmctx) -> bool;
 
+            enter_sync_call(vmctx: vmctx, caller_instance: u32, callee_async: u32, callee_instance: u32) -> bool;
+            exit_sync_call(vmctx: vmctx) -> bool;
+
             #[cfg(feature = "component-model-async")]
             backpressure_modify(vmctx: vmctx, caller_instance: u32, increment: u8) -> bool;
             #[cfg(feature = "component-model-async")]
@@ -185,8 +188,6 @@ 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

@@ -267,6 +267,7 @@ pub enum CoreDef {
     InstanceFlags(RuntimeComponentInstanceIndex),
     Trampoline(TrampolineIndex),
     UnsafeIntrinsic(ModuleInternedTypeIndex, UnsafeIntrinsic),
+    TaskMayBlock,
 
     /// This is a special variant not present in `info::CoreDef` which
     /// represents that this definition refers to a fused adapter function. This
@@ -475,8 +476,9 @@ pub enum Trampoline {
     FutureTransfer,
     StreamTransfer,
     ErrorContextTransfer,
-    CheckBlocking,
     Trap,
+    EnterSyncCall,
+    ExitSyncCall,
     ContextGet {
         instance: RuntimeComponentInstanceIndex,
         slot: u32,
@@ -905,6 +907,7 @@ impl LinearizeDfg<'_> {
                 }
                 info::CoreDef::UnsafeIntrinsic(*i)
             }
+            CoreDef::TaskMayBlock => info::CoreDef::TaskMayBlock,
         }
     }
 
@@ -1156,8 +1159,9 @@ impl LinearizeDfg<'_> {
             Trampoline::FutureTransfer => info::Trampoline::FutureTransfer,
             Trampoline::StreamTransfer => info::Trampoline::StreamTransfer,
             Trampoline::ErrorContextTransfer => info::Trampoline::ErrorContextTransfer,
-            Trampoline::CheckBlocking => info::Trampoline::CheckBlocking,
             Trampoline::Trap => info::Trampoline::Trap,
+            Trampoline::EnterSyncCall => info::Trampoline::EnterSyncCall,
+            Trampoline::ExitSyncCall => info::Trampoline::ExitSyncCall,
             Trampoline::ContextGet { instance, slot } => info::Trampoline::ContextGet {
                 instance: *instance,
                 slot: *slot,

crates/environ/src/component/info.rs

@@ -389,6 +389,10 @@ pub enum CoreDef {
     Trampoline(TrampolineIndex),
     /// An intrinsic for compile-time builtins.
     UnsafeIntrinsic(UnsafeIntrinsic),
+    /// Reference to a wasm global which represents a runtime-managed boolean
+    /// indicating whether the currently-running task may perform a blocking
+    /// operation.
+    TaskMayBlock,
 }
 
 impl<T> From<CoreExport<T>> for CoreDef
@@ -1105,9 +1109,13 @@ 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,
+    /// An intrinsic used by FACT-generated modules to check whether an instance
+    /// may be entered for a sync-to-sync call and push a task onto the stack if
+    /// so.
+    EnterSyncCall,
+    /// An intrinsic used by FACT-generated modules to pop the task previously
+    /// pushed by `EnterSyncCall`.
+    ExitSyncCall,
 
     /// An intrinsic used by FACT-generated modules to trap with a specified
     /// code.
@@ -1242,8 +1250,9 @@ impl Trampoline {
             FutureTransfer => format!("future-transfer"),
             StreamTransfer => format!("stream-transfer"),
             ErrorContextTransfer => format!("error-context-transfer"),
-            CheckBlocking => format!("check-blocking"),
             Trap => format!("trap"),
+            EnterSyncCall => format!("enter-sync-call"),
+            ExitSyncCall => format!("exit-sync-call"),
             ContextGet { .. } => format!("context-get"),
             ContextSet { .. } => format!("context-set"),
             ThreadIndex => format!("thread-index"),

crates/environ/src/component/translate.rs

@@ -79,6 +79,14 @@ pub struct Translator<'a, 'data> {
 
     /// The top-level import name for Wasmtime's unsafe intrinsics, if any.
     unsafe_intrinsics_import: Option<&'a str>,
+
+    /// Whether this component contains any module instantiations outside of
+    /// leaf components.
+    ///
+    /// We'll use this when generating fused adapters; if it's false, then we
+    /// know that no guest-to-guest call can reenter a runtime instance
+    /// recursively.
+    has_nonleaf_module_instantiations: bool,
 }
 
 /// Representation of the syntactic scope of a component meaning where it is
@@ -171,6 +179,14 @@ struct Translation<'data> {
     /// component has finished, e.g. for the `inline` pass, but beforehand this
     /// is set to `None`.
     types: Option<Types>,
+
+    /// Whether we've encountered a module instantiation while parsing this
+    /// component (disregarding instantiations in subcomponents).
+    saw_module_instantiation: bool,
+
+    /// Whether we've encountered a component instantation while parsing this
+    /// component (disregarding instantiations in subcomponents).
+    saw_component_instantiation: bool,
 }
 
 // NB: the type information contained in `LocalInitializer` should always point
@@ -455,6 +471,7 @@ impl<'a, 'data> Translator<'a, 'data> {
             static_modules: Default::default(),
             scope_vec,
             unsafe_intrinsics_import: None,
+            has_nonleaf_module_instantiations: false,
         }
     }
 
@@ -606,6 +623,7 @@ impl<'a, 'data> Translator<'a, 'data> {
 
                 let known_func = match arg {
                     CoreDef::InstanceFlags(_) => unreachable!("instance flags are not a function"),
+                    CoreDef::TaskMayBlock => unreachable!("task_may_block is not a function"),
 
                     // We could in theory inline these trampolines, so it could
                     // potentially make sense to record that we know this
@@ -694,6 +712,9 @@ impl<'a, 'data> Translator<'a, 'data> {
                 assert!(self.result.types.is_none());
                 self.result.types = Some(self.validator.end(offset)?);
 
+                self.has_nonleaf_module_instantiations |=
+                    self.result.saw_module_instantiation && self.result.saw_component_instantiation;
+
                 // Exit the current lexical scope. If there is no parent (no
                 // frame currently on the stack) then translation is finished.
                 // Otherwise that means that a nested component has been
@@ -1251,6 +1272,7 @@ impl<'a, 'data> Translator<'a, 'data> {
             // largely just records the arguments given from wasmparser into a
             // `HashMap` for processing later during inlining.
             Payload::InstanceSection(s) => {
+                self.result.saw_module_instantiation = true;
                 self.validator.instance_section(&s)?;
                 for instance in s {
                     let init = match instance? {
@@ -1266,6 +1288,7 @@ impl<'a, 'data> Translator<'a, 'data> {
                 }
             }
             Payload::ComponentInstanceSection(s) => {
+                self.result.saw_component_instantiation = true;
                 let mut index = self.validator.types(0).unwrap().component_instance_count();
                 self.validator.component_instance_section(&s)?;
                 for instance in s {

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

@@ -205,8 +205,11 @@ impl<'data> Translator<'_, 'data> {
         // the module using standard core wasm translation, and then fills out
         // the dfg metadata for each adapter.
         for (module_id, adapter_module) in state.adapter_modules.iter() {
-            let mut module =
-                fact::Module::new(self.types.types(), self.tunables.debug_adapter_modules);
+            let mut module = fact::Module::new(
+                self.types.types(),
+                self.tunables.debug_adapter_modules,
+                self.has_nonleaf_module_instantiations,
+            );
             let mut names = Vec::with_capacity(adapter_module.adapters.len());
             for adapter in adapter_module.adapters.iter() {
                 let name = format!("adapter{}", adapter.as_u32());
@@ -345,8 +348,9 @@ fn fact_import_to_core_def(
         fact::Import::ErrorContextTransfer => {
             simple_intrinsic(dfg::Trampoline::ErrorContextTransfer)
         }
-        fact::Import::CheckBlocking => simple_intrinsic(dfg::Trampoline::CheckBlocking),
         fact::Import::Trap => simple_intrinsic(dfg::Trampoline::Trap),
+        fact::Import::EnterSyncCall => simple_intrinsic(dfg::Trampoline::EnterSyncCall),
+        fact::Import::ExitSyncCall => simple_intrinsic(dfg::Trampoline::ExitSyncCall),
     }
 }
 
@@ -452,7 +456,8 @@ impl PartitionAdapterModules {
             // These items can't transitively depend on an adapter
             dfg::CoreDef::Trampoline(_)
             | dfg::CoreDef::InstanceFlags(_)
-            | dfg::CoreDef::UnsafeIntrinsic(..) => {}
+            | dfg::CoreDef::UnsafeIntrinsic(..)
+            | dfg::CoreDef::TaskMayBlock => {}
         }
     }
 

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

@@ -538,11 +538,9 @@ impl<'a> Inliner<'a> {
                     // happening within the same component instance.
                     //
                     // In this situation if the `canon.lower`'d function is
-                    // called then it immediately sets `may_enter` to `false`.
-                    // When calling the callee, however, that's `canon.lift`
-                    // which immediately traps if `may_enter` is `false`. That
-                    // means that this pairing of functions creates a function
-                    // that always traps.
+                    // called then it recursively re-enters itself, which is
+                    // defined to trap by the spec. That means that this pairing
+                    // of functions creates a function that always traps.
                     //
                     // When closely reading the spec though the precise trap
                     // that comes out can be somewhat variable. Technically the
@@ -1594,7 +1592,7 @@ impl<'a> Inliner<'a> {
         }
     }
 
-    /// Translatees an `AdapterOptions` into a `CanonicalOptions` where
+    /// Translates an `AdapterOptions` into a `CanonicalOptions` where
     /// memories/functions are inserted into the global initializer list for
     /// use at runtime. This is only used for lowered host functions and lifted
     /// functions exported to the host.

crates/environ/src/component/vmcomponent_offsets.rs

@@ -28,13 +28,9 @@ pub const VMCOMPONENT_MAGIC: u32 = u32::from_le_bytes(*b"comp");
 /// canonical ABI flag `may_leave`
 pub const FLAG_MAY_LEAVE: i32 = 1 << 0;
 
-/// Flag for the `VMComponentContext::flags` field which corresponds to the
-/// canonical ABI flag `may_enter`
-pub const FLAG_MAY_ENTER: i32 = 1 << 1;
-
 /// Flag for the `VMComponentContext::flags` field which is set whenever a
 /// function is called to indicate that `post_return` must be called next.
-pub const FLAG_NEEDS_POST_RETURN: i32 = 1 << 2;
+pub const FLAG_NEEDS_POST_RETURN: i32 = 1 << 1;
 
 /// Runtime offsets within a `VMComponentContext` for a specific component.
 #[derive(Debug, Clone, Copy)]
@@ -70,6 +66,7 @@ pub struct VMComponentOffsets<P> {
     builtins: u32,
     vm_store_context: u32,
     flags: u32,
+    task_may_block: u32,
     trampoline_func_refs: u32,
     intrinsic_func_refs: u32,
     lowerings: u32,
@@ -129,6 +126,7 @@ impl<P: PtrSize> VMComponentOffsets<P> {
             builtins: 0,
             vm_store_context: 0,
             flags: 0,
+            task_may_block: 0,
             trampoline_func_refs: 0,
             intrinsic_func_refs: 0,
             lowerings: 0,
@@ -172,6 +170,7 @@ impl<P: PtrSize> VMComponentOffsets<P> {
             size(vm_store_context) = ret.ptr.size(),
             align(16),
             size(flags) = cmul(ret.num_runtime_component_instances, ret.ptr.size_of_vmglobal_definition()),
+            size(task_may_block) = ret.ptr.size(),
             align(u32::from(ret.ptr.size())),
             size(trampoline_func_refs) = cmul(ret.num_trampolines, ret.ptr.size_of_vm_func_ref()),
             size(intrinsic_func_refs) = cmul(ret.num_unsafe_intrinsics, ret.ptr.size_of_vm_func_ref()),
@@ -219,6 +218,11 @@ impl<P: PtrSize> VMComponentOffsets<P> {
         self.flags + index.as_u32() * u32::from(self.ptr.size_of_vmglobal_definition())
     }
 
+    /// The offset of the `task_may_block` field.
+    pub fn task_may_block(&self) -> u32 {
+        self.task_may_block
+    }
+
     /// The offset of the `vm_store_context` field.
     #[inline]
     pub fn vm_store_context(&self) -> u32 {