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.

Author: dicej

crates/cranelift/src/compiler/component.rs

@@ -732,14 +732,6 @@ impl<'a> TrampolineCompiler<'a> {
                     |_, _| {},
                 );
             }
-            Trampoline::CheckBlocking => {
-                self.translate_libcall(
-                    host::check_blocking,
-                    TrapSentinel::Falsy,
-                    WasmArgs::InRegisters,
-                    |_, _| {},
-                );
-            }
             Trampoline::EnterSyncCall => {
                 self.translate_libcall(
                     host::enter_sync_call,

crates/environ/src/component.rs

@@ -188,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,7 +476,6 @@ pub enum Trampoline {
     FutureTransfer,
     StreamTransfer,
     ErrorContextTransfer,
-    CheckBlocking,
     EnterSyncCall,
     ExitSyncCall,
     ContextGet {
@@ -906,6 +906,7 @@ impl LinearizeDfg<'_> {
                 }
                 info::CoreDef::UnsafeIntrinsic(*i)
             }
+            CoreDef::TaskMayBlock => info::CoreDef::TaskMayBlock,
         }
     }
 
@@ -1157,7 +1158,6 @@ impl LinearizeDfg<'_> {
             Trampoline::FutureTransfer => info::Trampoline::FutureTransfer,
             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 {

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,10 +1109,6 @@ 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.
@@ -1246,7 +1246,6 @@ impl Trampoline {
             FutureTransfer => format!("future-transfer"),
             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"),

crates/environ/src/component/translate.rs

@@ -79,6 +79,13 @@ 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 modules 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_modules: bool,
 }
 
 /// Representation of the syntactic scope of a component meaning where it is
@@ -171,6 +178,13 @@ 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 while parsing this component
+    /// (but ignoring modules found in subcomponents).
+    saw_module: bool,
+
+    /// Whether we've encountered a subcomponent while parsing this component.
+    saw_component: bool,
 }
 
 // NB: the type information contained in `LocalInitializer` should always point
@@ -455,6 +469,7 @@ impl<'a, 'data> Translator<'a, 'data> {
             static_modules: Default::default(),
             scope_vec,
             unsafe_intrinsics_import: None,
+            has_nonleaf_modules: false,
         }
     }
 
@@ -606,6 +621,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 +710,8 @@ impl<'a, 'data> Translator<'a, 'data> {
                 assert!(self.result.types.is_none());
                 self.result.types = Some(self.validator.end(offset)?);
 
+                self.has_nonleaf_modules |= self.result.saw_module && self.result.saw_component;
+
                 // 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
@@ -1193,6 +1211,7 @@ impl<'a, 'data> Translator<'a, 'data> {
                 parser,
                 unchecked_range,
             } => {
+                self.result.saw_module = true;
                 let index = self.validator.types(0).unwrap().module_count();
                 self.validator.module_section(&unchecked_range)?;
                 let static_module_index = self.static_modules.next_key();
@@ -1238,6 +1257,7 @@ impl<'a, 'data> Translator<'a, 'data> {
                 parser,
                 unchecked_range,
             } => {
+                self.result.saw_component = true;
                 self.validator.component_section(&unchecked_range)?;
                 self.lexical_scopes.push(LexicalScope {
                     parser: mem::replace(&mut self.parser, parser),

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_modules,
+            );
             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,7 +348,6 @@ fn fact_import_to_core_def(
         fact::Import::ErrorContextTransfer => {
             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),
     }
@@ -453,7 +455,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

@@ -1592,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

@@ -66,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,
@@ -125,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,
@@ -168,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()),
@@ -215,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 {

crates/environ/src/fact.rs

@@ -54,6 +54,8 @@ pub static PREPARE_CALL_FIXED_PARAMS: &[ValType] = &[
 
 /// Representation of an adapter module.
 pub struct Module<'a> {
+    might_recursively_reenter: bool,
+
     /// Whether or not debug code is inserted into the adapters themselves.
     debug: bool,
     /// Type information from the creator of this `Module`
@@ -90,7 +92,6 @@ pub struct Module<'a> {
     imported_stream_transfer: Option<FuncIndex>,
     imported_error_context_transfer: Option<FuncIndex>,
 
-    imported_check_blocking: Option<FuncIndex>,
     imported_enter_sync_call: Option<FuncIndex>,
     imported_exit_sync_call: Option<FuncIndex>,
 
@@ -104,6 +105,8 @@ pub struct Module<'a> {
     helper_worklist: Vec<(FunctionId, Helper)>,
 
     exports: Vec<(u32, String)>,
+
+    task_may_block: Option<GlobalIndex>,
 }
 
 struct AdapterData {
@@ -241,8 +244,13 @@ enum HelperLocation {
 
 impl<'a> Module<'a> {
     /// Creates an empty module.
-    pub fn new(types: &'a ComponentTypesBuilder, debug: bool) -> Module<'a> {
+    pub fn new(
+        types: &'a ComponentTypesBuilder,
+        debug: bool,
+        might_recursively_reenter: bool,
+    ) -> Module<'a> {
         Module {
+            might_recursively_reenter,
             debug,
             types,
             core_types: Default::default(),
@@ -264,10 +272,10 @@ impl<'a> Module<'a> {
             imported_future_transfer: None,
             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(),
+            task_may_block: None,
         }
     }
 
@@ -463,6 +471,25 @@ impl<'a> Module<'a> {
         idx
     }
 
+    fn import_task_may_block(&mut self) -> GlobalIndex {
+        if let Some(task_may_block) = self.task_may_block {
+            task_may_block
+        } else {
+            let task_may_block = self.import_global(
+                "instance",
+                "task_may_block",
+                GlobalType {
+                    val_type: ValType::I32,
+                    mutable: true,
+                    shared: false,
+                },
+                CoreDef::TaskMayBlock,
+            );
+            self.task_may_block = Some(task_may_block);
+            task_may_block
+        }
+    }
+
     fn import_transcoder(&mut self, transcoder: transcode::Transcoder) -> FuncIndex {
         *self
             .imported_transcoders
@@ -720,17 +747,6 @@ 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 import_enter_sync_call(&mut self) -> FuncIndex {
         self.import_simple(
             "async",
@@ -911,9 +927,6 @@ 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,
     /// 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.