Refactor InstanceAllocator trait impl split (#11457)

Prior to this commit Wasmtime had an `InstanceAllocatorImpl` trait with
a number of required methods as well as an `InstanceAllocator` trait
with a number of provided impls. The `InstanceAllocator` trait is
implemented for everything implementing `InstanceAllocatorImpl` to force
users to be unable to override the default methods. When adding `async`
support internally to Wasmtime these are going to need to be
`#[async_trait]`-annotated-traits which adds a cost to `async` functions
as a future needs to be heap-allocated.

The goal of this commit is to make this future `async`-ification a bit
more optimal. Notably the `InstanceAllocator` trait is removed and
replaced with inherent methods on `impl dyn InstanceAllocatorImpl`.
After that the previous `InstanceAllocatorImpl` trait was renamed to
`InstanceAllocator` meaning that there's just one `InstanceAllocator`
trait which has inherent methods which cannot be overridden. A
consequence of this is that the inherent methods are also forced to do
virtual dispatch unlike before where they would internally use
monomorphization to have static dispatch. Given the complexity of
instance allocation this is expected to be a negligible cost, however.

The main benefit is that `allocate_module`, `allocate_tables`, and
`allocate_memories` all get to be native `async` functions without the
cost of `#[async_trait]`. Only allocation of a single table/memory will
require an allocation of a future which in profiling helps reduce the
cost of instantiation slightly.

Note that `#[async_trait]` is not currently used, this commit is just
preparation for its eventual use.

Author: alexcrichton

crates/wasmtime/src/runtime/store.rs

@@ -2505,11 +2505,11 @@ impl Drop for StoreOpaque {
 
             for (id, instance) in self.instances.iter_mut() {
                 log::trace!("store {store_id:?} is deallocating {id:?}");
-                if let StoreInstanceKind::Dummy = instance.kind {
-                    ondemand.deallocate_module(&mut instance.handle);
-                } else {
-                    allocator.deallocate_module(&mut instance.handle);
-                }
+                let allocator = match instance.kind {
+                    StoreInstanceKind::Dummy => &ondemand,
+                    _ => allocator,
+                };
+                allocator.deallocate_module(&mut instance.handle);
             }
 
             #[cfg(feature = "component-model")]

crates/wasmtime/src/runtime/trampoline/memory.rs

@@ -3,9 +3,9 @@ use crate::memory::{LinearMemory, MemoryCreator};
 use crate::prelude::*;
 use crate::runtime::vm::mpk::ProtectionKey;
 use crate::runtime::vm::{
-    CompiledModuleId, InstanceAllocationRequest, InstanceAllocatorImpl, Memory,
-    MemoryAllocationIndex, MemoryBase, ModuleRuntimeInfo, OnDemandInstanceAllocator,
-    RuntimeLinearMemory, RuntimeMemoryCreator, SharedMemory, Table, TableAllocationIndex,
+    CompiledModuleId, InstanceAllocationRequest, InstanceAllocator, Memory, MemoryAllocationIndex,
+    MemoryBase, ModuleRuntimeInfo, OnDemandInstanceAllocator, RuntimeLinearMemory,
+    RuntimeMemoryCreator, SharedMemory, Table, TableAllocationIndex,
 };
 use crate::store::{AllocateInstanceKind, InstanceId, StoreOpaque};
 use alloc::sync::Arc;
@@ -122,9 +122,9 @@ struct SingleMemoryInstance<'a> {
     ondemand: OnDemandInstanceAllocator,
 }
 
-unsafe impl InstanceAllocatorImpl for SingleMemoryInstance<'_> {
+unsafe impl InstanceAllocator for SingleMemoryInstance<'_> {
     #[cfg(feature = "component-model")]
-    fn validate_component_impl<'a>(
+    fn validate_component<'a>(
         &self,
         _component: &Component,
         _offsets: &VMComponentOffsets<HostPtr>,
@@ -133,18 +133,18 @@ unsafe impl InstanceAllocatorImpl for SingleMemoryInstance<'_> {
         unreachable!("`SingleMemoryInstance` allocator never used with components")
     }
 
-    fn validate_module_impl(&self, module: &Module, offsets: &VMOffsets<HostPtr>) -> Result<()> {
+    fn validate_module(&self, module: &Module, offsets: &VMOffsets<HostPtr>) -> Result<()> {
         anyhow::ensure!(
             module.memories.len() == 1,
             "`SingleMemoryInstance` allocator can only be used for modules with a single memory"
         );
-        self.ondemand.validate_module_impl(module, offsets)?;
+        self.ondemand.validate_module(module, offsets)?;
         Ok(())
     }
 
     #[cfg(feature = "gc")]
-    fn validate_memory_impl(&self, memory: &wasmtime_environ::Memory) -> Result<()> {
-        self.ondemand.validate_memory_impl(memory)
+    fn validate_memory(&self, memory: &wasmtime_environ::Memory) -> Result<()> {
+        self.ondemand.validate_memory(memory)
     }
 
     #[cfg(feature = "component-model")]
@@ -172,11 +172,10 @@ unsafe impl InstanceAllocatorImpl for SingleMemoryInstance<'_> {
         tunables: &Tunables,
         memory_index: Option<DefinedMemoryIndex>,
     ) -> Result<(MemoryAllocationIndex, Memory)> {
-        #[cfg(debug_assertions)]
-        {
+        if cfg!(debug_assertions) {
             let module = request.runtime_info.env_module();
             let offsets = request.runtime_info.offsets();
-            self.validate_module_impl(module, offsets)
+            self.validate_module(module, offsets)
                 .expect("should have already validated the module before allocating memory");
         }
 

crates/wasmtime/src/runtime/vm.rs

@@ -98,9 +98,9 @@ pub use crate::runtime::vm::export::*;
 pub use crate::runtime::vm::gc::*;
 pub use crate::runtime::vm::imports::Imports;
 pub use crate::runtime::vm::instance::{
-    GcHeapAllocationIndex, Instance, InstanceAllocationRequest, InstanceAllocator,
-    InstanceAllocatorImpl, InstanceHandle, MemoryAllocationIndex, OnDemandInstanceAllocator,
-    StorePtr, TableAllocationIndex, initialize_instance,
+    GcHeapAllocationIndex, Instance, InstanceAllocationRequest, InstanceAllocator, InstanceHandle,
+    MemoryAllocationIndex, OnDemandInstanceAllocator, StorePtr, TableAllocationIndex,
+    initialize_instance,
 };
 #[cfg(feature = "pooling-allocator")]
 pub use crate::runtime::vm::instance::{

crates/wasmtime/src/runtime/vm/instance/allocator.rs

@@ -197,23 +197,23 @@ impl GcHeapAllocationIndex {
 ///
 /// This trait is unsafe as it requires knowledge of Wasmtime's runtime
 /// internals to implement correctly.
-pub unsafe trait InstanceAllocatorImpl {
+pub unsafe trait InstanceAllocator: Send + Sync {
     /// Validate whether a component (including all of its contained core
     /// modules) is allocatable by this instance allocator.
     #[cfg(feature = "component-model")]
-    fn validate_component_impl<'a>(
+    fn validate_component<'a>(
         &self,
         component: &Component,
         offsets: &VMComponentOffsets<HostPtr>,
         get_module: &'a dyn Fn(StaticModuleIndex) -> &'a Module,
     ) -> Result<()>;
 
     /// Validate whether a module is allocatable by this instance allocator.
-    fn validate_module_impl(&self, module: &Module, offsets: &VMOffsets<HostPtr>) -> Result<()>;
+    fn validate_module(&self, module: &Module, offsets: &VMOffsets<HostPtr>) -> Result<()>;
 
     /// Validate whether a memory is allocatable by this instance allocator.
     #[cfg(feature = "gc")]
-    fn validate_memory_impl(&self, memory: &wasmtime_environ::Memory) -> Result<()>;
+    fn validate_memory(&self, memory: &wasmtime_environ::Memory) -> Result<()>;
 
     /// Increment the count of concurrent component instances that are currently
     /// allocated, if applicable.
@@ -230,7 +230,7 @@ pub unsafe trait InstanceAllocatorImpl {
     //    associated types are not object safe.
     //
     // 2. We would want a parameterized `Drop` implementation so that we could
-    //    pass in the `InstanceAllocatorImpl` on drop, but this doesn't exist in
+    //    pass in the `InstanceAllocator` on drop, but this doesn't exist in
     //    Rust. Therefore, we would be forced to add reference counting and
     //    stuff like that to keep a handle on the instance allocator from this
     //    theoretical type. That's a bummer.
@@ -359,36 +359,7 @@ pub unsafe trait InstanceAllocatorImpl {
     fn allow_all_pkeys(&self);
 }
 
-/// A thing that can allocate instances.
-///
-/// Don't implement this trait directly, instead implement
-/// `InstanceAllocatorImpl` and you'll get this trait for free via a blanket
-/// impl.
-pub trait InstanceAllocator: InstanceAllocatorImpl {
-    /// Validate whether a component (including all of its contained core
-    /// modules) is allocatable with this instance allocator.
-    #[cfg(feature = "component-model")]
-    fn validate_component<'a>(
-        &self,
-        component: &Component,
-        offsets: &VMComponentOffsets<HostPtr>,
-        get_module: &'a dyn Fn(StaticModuleIndex) -> &'a Module,
-    ) -> Result<()> {
-        InstanceAllocatorImpl::validate_component_impl(self, component, offsets, get_module)
-    }
-
-    /// Validate whether a core module is allocatable with this instance
-    /// allocator.
-    fn validate_module(&self, module: &Module, offsets: &VMOffsets<HostPtr>) -> Result<()> {
-        InstanceAllocatorImpl::validate_module_impl(self, module, offsets)
-    }
-
-    /// Validate whether a memory is allocatable with this instance allocator.
-    #[cfg(feature = "gc")]
-    fn validate_memory(&self, memory: &wasmtime_environ::Memory) -> Result<()> {
-        InstanceAllocatorImpl::validate_memory_impl(self, memory)
-    }
-
+impl dyn InstanceAllocator + '_ {
     /// Allocates a fresh `InstanceHandle` for the `req` given.
     ///
     /// This will allocate memories and tables internally from this allocator
@@ -402,15 +373,16 @@ pub trait InstanceAllocator: InstanceAllocatorImpl {
     ///
     /// The `request` provided must be valid, e.g. the imports within are
     /// correctly sized/typed for the instance being created.
-    unsafe fn allocate_module(
+    pub(crate) unsafe fn allocate_module(
         &self,
         mut request: InstanceAllocationRequest,
     ) -> Result<InstanceHandle> {
         let module = request.runtime_info.env_module();
 
-        #[cfg(debug_assertions)]
-        InstanceAllocatorImpl::validate_module_impl(self, module, request.runtime_info.offsets())
-            .expect("module should have already been validated before allocation");
+        if cfg!(debug_assertions) {
+            InstanceAllocator::validate_module(self, module, request.runtime_info.offsets())
+                .expect("module should have already been validated before allocation");
+        }
 
         self.increment_core_instance_count()?;
 
@@ -449,7 +421,7 @@ pub trait InstanceAllocator: InstanceAllocatorImpl {
     /// # Unsafety
     ///
     /// The instance must have previously been allocated by `Self::allocate`.
-    unsafe fn deallocate_module(&self, handle: &mut InstanceHandle) {
+    pub(crate) unsafe fn deallocate_module(&self, handle: &mut InstanceHandle) {
         // SAFETY: the contract of `deallocate_*` is itself a contract of this
         // function, that the memories/tables were previously allocated from
         // here.
@@ -470,9 +442,10 @@ pub trait InstanceAllocator: InstanceAllocatorImpl {
     ) -> Result<()> {
         let module = request.runtime_info.env_module();
 
-        #[cfg(debug_assertions)]
-        InstanceAllocatorImpl::validate_module_impl(self, module, request.runtime_info.offsets())
-            .expect("module should have already been validated before allocation");
+        if cfg!(debug_assertions) {
+            InstanceAllocator::validate_module(self, module, request.runtime_info.offsets())
+                .expect("module should have already been validated before allocation");
+        }
 
         for (memory_index, ty) in module.memories.iter().skip(module.num_imported_memories) {
             let memory_index = module
@@ -520,9 +493,10 @@ pub trait InstanceAllocator: InstanceAllocatorImpl {
     ) -> Result<()> {
         let module = request.runtime_info.env_module();
 
-        #[cfg(debug_assertions)]
-        InstanceAllocatorImpl::validate_module_impl(self, module, request.runtime_info.offsets())
-            .expect("module should have already been validated before allocation");
+        if cfg!(debug_assertions) {
+            InstanceAllocator::validate_module(self, module, request.runtime_info.offsets())
+                .expect("module should have already been validated before allocation");
+        }
 
         for (index, table) in module.tables.iter().skip(module.num_imported_tables) {
             let def_index = module
@@ -556,11 +530,6 @@ pub trait InstanceAllocator: InstanceAllocatorImpl {
     }
 }
 
-// Every `InstanceAllocatorImpl` is an `InstanceAllocator` when used
-// correctly. Also, no one is allowed to override this trait's methods, they
-// must use the defaults. This blanket impl provides both of those things.
-impl<T: InstanceAllocatorImpl> InstanceAllocator for T {}
-
 fn check_table_init_bounds(
     store: &mut StoreOpaque,
     instance: InstanceId,
@@ -894,7 +863,6 @@ mod tests {
 
     #[test]
     fn allocator_traits_are_object_safe() {
-        fn _instance_allocator(_: &dyn InstanceAllocatorImpl) {}
-        fn _instance_allocator_ext(_: &dyn InstanceAllocator) {}
+        fn _instance_allocator(_: &dyn InstanceAllocator) {}
     }
 }

crates/wasmtime/src/runtime/vm/instance/allocator/on_demand.rs

@@ -1,5 +1,5 @@
 use super::{
-    InstanceAllocationRequest, InstanceAllocatorImpl, MemoryAllocationIndex, TableAllocationIndex,
+    InstanceAllocationRequest, InstanceAllocator, MemoryAllocationIndex, TableAllocationIndex,
 };
 use crate::prelude::*;
 use crate::runtime::vm::CompiledModuleId;
@@ -76,9 +76,9 @@ impl Default for OnDemandInstanceAllocator {
     }
 }
 
-unsafe impl InstanceAllocatorImpl for OnDemandInstanceAllocator {
+unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
     #[cfg(feature = "component-model")]
-    fn validate_component_impl<'a>(
+    fn validate_component<'a>(
         &self,
         _component: &Component,
         _offsets: &VMComponentOffsets<HostPtr>,
@@ -87,12 +87,12 @@ unsafe impl InstanceAllocatorImpl for OnDemandInstanceAllocator {
         Ok(())
     }
 
-    fn validate_module_impl(&self, _module: &Module, _offsets: &VMOffsets<HostPtr>) -> Result<()> {
+    fn validate_module(&self, _module: &Module, _offsets: &VMOffsets<HostPtr>) -> Result<()> {
         Ok(())
     }
 
     #[cfg(feature = "gc")]
-    fn validate_memory_impl(&self, _memory: &wasmtime_environ::Memory) -> Result<()> {
+    fn validate_memory(&self, _memory: &wasmtime_environ::Memory) -> Result<()> {
         Ok(())
     }
 

crates/wasmtime/src/runtime/vm/instance/allocator/pooling.rs

@@ -44,7 +44,7 @@ use self::decommit_queue::DecommitQueue;
 use self::memory_pool::MemoryPool;
 use self::table_pool::TablePool;
 use super::{
-    InstanceAllocationRequest, InstanceAllocatorImpl, MemoryAllocationIndex, TableAllocationIndex,
+    InstanceAllocationRequest, InstanceAllocator, MemoryAllocationIndex, TableAllocationIndex,
 };
 use crate::Enabled;
 use crate::prelude::*;
@@ -523,9 +523,9 @@ impl PoolingInstanceAllocator {
     }
 }
 
-unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
+unsafe impl InstanceAllocator for PoolingInstanceAllocator {
     #[cfg(feature = "component-model")]
-    fn validate_component_impl<'a>(
+    fn validate_component<'a>(
         &self,
         component: &Component,
         offsets: &VMComponentOffsets<HostPtr>,
@@ -549,7 +549,7 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
                 InstantiateModule(InstantiateModule::Static(static_module_index, _)) => {
                     let module = get_module(*static_module_index);
                     let offsets = VMOffsets::new(HostPtr, &module);
-                    self.validate_module_impl(module, &offsets)?;
+                    self.validate_module(module, &offsets)?;
                     num_core_instances += 1;
                     num_memories += module.num_defined_memories();
                     num_tables += module.num_defined_tables();
@@ -593,7 +593,7 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
         Ok(())
     }
 
-    fn validate_module_impl(&self, module: &Module, offsets: &VMOffsets<HostPtr>) -> Result<()> {
+    fn validate_module(&self, module: &Module, offsets: &VMOffsets<HostPtr>) -> Result<()> {
         self.validate_memory_plans(module)
             .context("module memory does not fit in pooling allocator requirements")?;
         self.validate_table_plans(module)
@@ -604,7 +604,7 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
     }
 
     #[cfg(feature = "gc")]
-    fn validate_memory_impl(&self, memory: &wasmtime_environ::Memory) -> Result<()> {
+    fn validate_memory(&self, memory: &wasmtime_environ::Memory) -> Result<()> {
         self.memories.validate_memory(memory)
     }