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.

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

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 anyhow::{Result, bail};
 use cranelift_codegen::ir::condcodes::IntCC;
 use cranelift_codegen::ir::{self, InstBuilder, MemFlags, Value};
@@ -740,6 +740,22 @@ impl<'a> TrampolineCompiler<'a> {
                     |_, _| {},
                 );
             }
+            Trampoline::EnterSyncCall => {
+                self.translate_libcall(
+                    host::enter_sync_call,
+                    TrapSentinel::Falsy,
+                    WasmArgs::InRegisters,
+                    |_, _| {},
+                );
+            }
+            Trampoline::ExitSyncCall => {
+                self.translate_libcall(
+                    host::exit_sync_call,
+                    TrapSentinel::Falsy,
+                    WasmArgs::InRegisters,
+                    |_, _| {},
+                );
+            }
             Trampoline::ContextGet { instance, slot } => {
                 self.translate_libcall(
                     host::context_get,
@@ -1129,11 +1145,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
@@ -1147,6 +1161,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:
@@ -1177,25 +1194,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.
@@ -1244,6 +1268,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")]

crates/environ/src/component/dfg.rs

@@ -476,6 +476,8 @@ pub enum Trampoline {
     StreamTransfer,
     ErrorContextTransfer,
     CheckBlocking,
+    EnterSyncCall,
+    ExitSyncCall,
     ContextGet {
         instance: RuntimeComponentInstanceIndex,
         slot: u32,
@@ -1156,6 +1158,8 @@ impl LinearizeDfg<'_> {
             Trampoline::StreamTransfer => info::Trampoline::StreamTransfer,
             Trampoline::ErrorContextTransfer => info::Trampoline::ErrorContextTransfer,
             Trampoline::CheckBlocking => info::Trampoline::CheckBlocking,
+            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

@@ -1109,6 +1109,14 @@ pub enum Trampoline {
     /// 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,
+
     /// Intrinsic used to implement the `context.get` component model builtin.
     ///
     /// The payload here represents that this is accessing the Nth slot of local
@@ -1239,6 +1247,8 @@ impl Trampoline {
             StreamTransfer => format!("stream-transfer"),
             ErrorContextTransfer => format!("error-context-transfer"),
             CheckBlocking => format!("check-blocking"),
+            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/adapt.rs

@@ -346,6 +346,8 @@ fn fact_import_to_core_def(
             simple_intrinsic(dfg::Trampoline::ErrorContextTransfer)
         }
         fact::Import::CheckBlocking => simple_intrinsic(dfg::Trampoline::CheckBlocking),
+        fact::Import::EnterSyncCall => simple_intrinsic(dfg::Trampoline::EnterSyncCall),
+        fact::Import::ExitSyncCall => simple_intrinsic(dfg::Trampoline::ExitSyncCall),
     }
 }
 

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

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

crates/environ/src/fact.rs

@@ -91,6 +91,8 @@ pub struct Module<'a> {
     imported_error_context_transfer: Option<FuncIndex>,
 
     imported_check_blocking: Option<FuncIndex>,
+    imported_enter_sync_call: Option<FuncIndex>,
+    imported_exit_sync_call: Option<FuncIndex>,
 
     // Current status of index spaces from the imports generated so far.
     imported_funcs: PrimaryMap<FuncIndex, Option<CoreDef>>,
@@ -263,6 +265,8 @@ impl<'a> Module<'a> {
             imported_stream_transfer: None,
             imported_error_context_transfer: None,
             imported_check_blocking: None,
+            imported_enter_sync_call: None,
+            imported_exit_sync_call: None,
             exports: Vec::new(),
         }
     }
@@ -727,6 +731,28 @@ impl<'a> Module<'a> {
         )
     }
 
+    fn import_enter_sync_call(&mut self) -> FuncIndex {
+        self.import_simple(
+            "async",
+            "enter-sync-call",
+            &[ValType::I32; 3],
+            &[],
+            Import::EnterSyncCall,
+            |me| &mut me.imported_enter_sync_call,
+        )
+    }
+
+    fn import_exit_sync_call(&mut self) -> FuncIndex {
+        self.import_simple(
+            "async",
+            "exit-sync-call",
+            &[],
+            &[],
+            Import::ExitSyncCall,
+            |me| &mut me.imported_exit_sync_call,
+        )
+    }
+
     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
@@ -888,6 +914,13 @@ pub enum Import {
     /// 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,
 }
 
 impl Options {

crates/environ/src/fact/trampoline.rs

@@ -16,8 +16,8 @@
 //! can be somewhat arbitrary, an intentional decision.
 
 use crate::component::{
-    CanonicalAbiInfo, ComponentTypesBuilder, FLAG_MAY_ENTER, FLAG_MAY_LEAVE, FixedEncoding as FE,
-    FlatType, InterfaceType, MAX_FLAT_ASYNC_PARAMS, MAX_FLAT_PARAMS, PREPARE_ASYNC_NO_RESULT,
+    CanonicalAbiInfo, ComponentTypesBuilder, FLAG_MAY_LEAVE, FixedEncoding as FE, FlatType,
+    InterfaceType, MAX_FLAT_ASYNC_PARAMS, MAX_FLAT_PARAMS, PREPARE_ASYNC_NO_RESULT,
     PREPARE_ASYNC_WITH_RESULT, START_FLAG_ASYNC_CALLEE, StringEncoding, Transcode,
     TypeComponentLocalErrorContextTableIndex, TypeEnumIndex, TypeFlagsIndex, TypeFutureTableIndex,
     TypeListIndex, TypeOptionIndex, TypeRecordIndex, TypeResourceTableIndex, TypeResultIndex,
@@ -743,17 +743,42 @@ impl<'a, 'b> Compiler<'a, 'b> {
         // flags on the callee if necessary whether it can be entered.
         self.trap_if_not_flag(adapter.lower.flags, FLAG_MAY_LEAVE, Trap::CannotLeave);
         if adapter.called_as_export {
-            self.trap_if_not_flag(adapter.lift.flags, FLAG_MAY_ENTER, Trap::CannotEnter);
-            self.set_flag(adapter.lift.flags, FLAG_MAY_ENTER, false);
-        } else if self.module.debug {
-            self.assert_not_flag(
-                adapter.lift.flags,
-                FLAG_MAY_ENTER,
-                "may_enter should be unset",
-            );
-        }
-
-        if self.types[adapter.lift.ty].async_ {
+            // Here we call a runtime intrinsic to check whether the instance
+            // we're about to enter has already been entered by this task
+            // (i.e. whether this is a recursive reentrance).  If so, we'll trap
+            // per the component model specification.  Otherwise, we'll push a
+            // task onto the stack and then later pop it off after the call.
+            //
+            // Note that, prior to the introduction of concurrency to the
+            // component model, we were able to handle this via a per-instance
+            // variable, allowing us to check and set that variable without
+            // leaving Wasm code.  However, now that there may be more than one
+            // concurrent task running in a given instance, a per-instance
+            // variable is not sufficient.
+            //
+            // Fortunately, we still have a number of opportunities to optimize
+            // common cases (e.g. components which contain only sync functions,
+            // and compositions of fewer than 64 transitive subcomponents, for
+            // which a per-task `uint64` bitset would suffice) and avoid the
+            // host call.  See the discussion of `trap_if_on_the_stack` in
+            // https://github.com/WebAssembly/component-model/blob/main/design/mvp/CanonicalABI.md
+            // for details.
+            //
+            // TODO: Implement one or more of those optimizations.
+            self.instruction(I32Const(
+                i32::try_from(adapter.lower.instance.as_u32()).unwrap(),
+            ));
+            self.instruction(I32Const(if self.types[adapter.lift.ty].async_ {
+                1
+            } else {
+                0
+            }));
+            self.instruction(I32Const(
+                i32::try_from(adapter.lift.instance.as_u32()).unwrap(),
+            ));
+            let enter_sync_call = self.module.import_enter_sync_call();
+            self.instruction(Call(enter_sync_call.as_u32()));
+        } else if self.types[adapter.lift.ty].async_ {
             let check_blocking = self.module.import_check_blocking();
             self.instruction(Call(check_blocking.as_u32()));
         }
@@ -818,9 +843,6 @@ impl<'a, 'b> Compiler<'a, 'b> {
             }
             self.instruction(Call(func.as_u32()));
         }
-        if adapter.called_as_export {
-            self.set_flag(adapter.lift.flags, FLAG_MAY_ENTER, true);
-        }
 
         for tmp in temps {
             self.free_temp_local(tmp);
@@ -831,6 +853,11 @@ impl<'a, 'b> Compiler<'a, 'b> {
             self.instruction(Call(exit.as_u32()));
         }
 
+        if adapter.called_as_export {
+            let exit_sync_call = self.module.import_exit_sync_call();
+            self.instruction(Call(exit_sync_call.as_u32()));
+        }
+
         self.finish()
     }
 
@@ -3268,15 +3295,6 @@ impl<'a, 'b> Compiler<'a, 'b> {
         self.instruction(End);
     }
 
-    fn assert_not_flag(&mut self, flags_global: GlobalIndex, flag_to_test: i32, msg: &'static str) {
-        self.instruction(GlobalGet(flags_global.as_u32()));
-        self.instruction(I32Const(flag_to_test));
-        self.instruction(I32And);
-        self.instruction(If(BlockType::Empty));
-        self.trap(Trap::AssertFailed(msg));
-        self.instruction(End);
-    }
-
     fn set_flag(&mut self, flags_global: GlobalIndex, flag_to_set: i32, value: bool) {
         self.instruction(GlobalGet(flags_global.as_u32()));
         if value {

crates/environ/src/fact/traps.rs

@@ -26,7 +26,6 @@ use wasm_encoder::Encode;
 #[derive(Hash, PartialEq, Eq, Copy, Clone)]
 pub enum Trap {
     CannotLeave,
-    CannotEnter,
     UnalignedPointer,
     InvalidDiscriminant,
     InvalidChar,
@@ -103,7 +102,6 @@ impl fmt::Display for Trap {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             Trap::CannotLeave => "cannot leave instance".fmt(f),
-            Trap::CannotEnter => "cannot enter instance".fmt(f),
             Trap::UnalignedPointer => "pointer not aligned correctly".fmt(f),
             Trap::InvalidDiscriminant => "invalid variant discriminant".fmt(f),
             Trap::InvalidChar => "invalid char value specified".fmt(f),