re-implement recursive reentrance checks

Now that the component model supports multiple tasks and cooperative threads
running concurrently in the same instance, the old model of using a per-instance
flag to track whether a component may be entered no longer works.  Instead of
keeping track of whether an instance has been entered at all, we must now keep
track of whether an instance has been entered for each in-progress task.

This commit removes `FLAG_MAY_ENTER` from `InstanceFlags`, relying instead on a
per-task call stack maintained at runtime.  When the `component-model-async`
feature is enabled, we reuse the `GuestTask` stack for this purpose.  When the
`component-model-async` feature is disabled, there's just a single stack used
for the entire store.  Note that these stacks are only updated when crossing
component boundaries -- not at every internal call within an instance.

Each time we're about to enter a instance from either the host or the guest, we
check the call stack, and if the instance is already present, we trap.
Otherwise, we push a new element onto the stack and later pop it back off once
the callee returns a value.

In addition to trapping on recursive reentrance, we do a couple of additional
checks for host-to-guest calls, for which we previously relied on
`FLAG_MAY_ENTER`:

- If _any_ instance owned by the store has previously trapped, we disallow entering that or any other instance owned by the store.

- If a post-return call is needed for an instance, we disallow entering it until that call has been made.

Note that, for sync-to-sync, guest-to-guest calls, the above process entails
significantly more overhead than the prior code, which only involved checking
and setting a flag and did not require calling into the host at all.  I intend
to follow this PR up with one or more optimizations to reduce that overhead.
See the discussion of `trap_if_on_the_stack` in
https://github.com/WebAssembly/component-model/blob/main/design/mvp/CanonicalABI.md
for examples of such optimizations.

In addition to the above refactoring, this commit does a couple of related
things:

- When a host function recursively calls back into guest code, we now chain the old call stack to the new one.  This allows us to catch recursive reentrance which may span multiple top-level instances.

- Previously, we did not push a new task on the stack for sync-to-sync, guest-to-guest calls, which meant we missed catching some violations of the sync-task-must-not-block-before-returning rule.  Now that we are pushing a new task in that case, we catch all such violations, which means some of the existing WAST tests needed updating.

bless disas/component-model tests

optimize sync-to-sync, guest-to-guest calls

We now omit recursive reentrance checks entirely when generating fused adapters
for components which cannot possibly recursively reenter themselves, noting that
components which only contain modules in leaf (sub)components fall into that
category.

However, we must still include a runtime global variable check-and-set to
enforce the may-block rules for sync-typed tasks.  With some more effort, we
could also eliminate those for components statically known to never make
blocking calls to intrinsics or imports.

The implementation of the may-block check centers around a
per-top-level-instance global variable called `task_may_block`, which we update
each time we switch threads and tasks, as well as whenever `task.return` or
`task.cancel` is called.  This required shuffling some code around and creating
a new `StoreOpaque::set_thread` function which encapsulates switching threads
and updating `task_may_block` for the old and new instances.

tweak rules to determine whether recursive reentrance is possible

This accounts for components which instantiate imported modules and/or
components.

Author: dicej

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 {