Lazily populate a store's trampoline map (#3742)

* Lazily populate a store's trampoline map

This commit is another installment of "how fast can we make
instantiation". Currently when instantiating a module with many function
imports each function, typically from the host, is inserted into the
store. This insertion process stores the `VMTrampoline` for the host
function in a side table so it can be looked up later if the host
function is called through the `Func` interface. This insertion process,
however, involves a hash map insertion which can be relatively expensive
at the scale of the rest of the instantiation process.

The optimization implemented in this commit is to avoid inserting
trampolines into the store at `Func`-insertion-time (aka instantiation
time) and instead only lazily populate the map of trampolines when
needed. The theory behind this is that almost all `Func` instances that
are called indirectly from the host are actually wasm functions, not
host-defined functions. This means that they already don't need to go
through the map of host trampolines and can instead be looked up from
the module they're defined in. With the assumed rarity of host functions
making `lookup_trampoline` a bit slower seems ok.

The `lookup_trampoline` function will now, on a miss from the wasm
modules and `host_trampolines` map, lazily iterate over the functions
within the store and insert trampolines into the `host_trampolines` map.
This process will eventually reach something which matches the function
provided because it should at least hit the same host function. The
relevant `lookup_trampoline` now sports a new documentation block
explaining all this as well for future readers.

Concretely this commit speeds up instantiation of an empty module with
100 imports and ~80 unique signatures from 10.6us to 6.4us, a 40%
improvement.

* Review comments

* Remove debug assert

crates/wasmtime/src/func.rs

@@ -718,7 +718,7 @@ impl Func {
     }
 
     pub(crate) fn sig_index(&self, data: &StoreData) -> VMSharedSignatureIndex {
-        unsafe { data[self.0].export().anyfunc.as_ref().type_index }
+        data[self.0].sig_index()
     }
 
     /// Invokes this function with the `params` given and writes returned values
@@ -2083,18 +2083,18 @@ impl HostFunc {
     /// Can only be inserted into stores with a matching `T` relative to when
     /// this `HostFunc` was first created.
     pub unsafe fn to_func(self: &Arc<Self>, store: &mut StoreOpaque) -> Func {
-        self.register_trampoline(store);
+        self.validate_store(store);
         let me = self.clone();
         Func::from_func_kind(FuncKind::SharedHost(me), store)
     }
 
     /// Same as [`HostFunc::to_func`], different ownership.
     unsafe fn into_func(self, store: &mut StoreOpaque) -> Func {
-        self.register_trampoline(store);
+        self.validate_store(store);
         Func::from_func_kind(FuncKind::Host(self), store)
     }
 
-    unsafe fn register_trampoline(&self, store: &mut StoreOpaque) {
+    fn validate_store(&self, store: &mut StoreOpaque) {
         // This assert is required to ensure that we can indeed safely insert
         // `self` into the `store` provided, otherwise the type information we
         // have listed won't be correct. This is possible to hit with the public
@@ -2103,8 +2103,6 @@ impl HostFunc {
             Engine::same(&self.engine, store.engine()),
             "cannot use a store with a different engine than a linker was created with",
         );
-        let idx = self.export.anyfunc.as_ref().type_index;
-        store.register_host_trampoline(idx, self.trampoline);
     }
 
     pub(crate) fn sig_index(&self) -> VMSharedSignatureIndex {
@@ -2128,7 +2126,7 @@ impl Drop for HostFunc {
 
 impl FuncData {
     #[inline]
-    fn trampoline(&self) -> VMTrampoline {
+    pub(crate) fn trampoline(&self) -> VMTrampoline {
         match &self.kind {
             FuncKind::StoreOwned { trampoline, .. } => *trampoline,
             FuncKind::SharedHost(host) => host.trampoline,
@@ -2140,6 +2138,10 @@ impl FuncData {
     fn export(&self) -> &ExportFunction {
         self.kind.export()
     }
+
+    pub(crate) fn sig_index(&self) -> VMSharedSignatureIndex {
+        unsafe { self.export().anyfunc.as_ref().type_index }
+    }
 }
 
 impl FuncKind {

crates/wasmtime/src/store.rs

@@ -259,7 +259,12 @@ pub struct StoreOpaque {
     signal_handler: Option<Box<SignalHandler<'static>>>,
     externref_activations_table: VMExternRefActivationsTable,
     modules: ModuleRegistry,
+
+    // See documentation on `StoreOpaque::lookup_trampoline` for what these
+    // fields are doing.
     host_trampolines: HashMap<VMSharedSignatureIndex, VMTrampoline>,
+    host_func_trampolines_registered: usize,
+
     // Numbers of resources instantiated in this store, and their limits
     instance_count: usize,
     instance_limit: usize,
@@ -436,6 +441,7 @@ impl<T> Store<T> {
                 externref_activations_table: VMExternRefActivationsTable::new(),
                 modules: ModuleRegistry::default(),
                 host_trampolines: HashMap::default(),
+                host_func_trampolines_registered: 0,
                 instance_count: 0,
                 instance_limit: crate::DEFAULT_INSTANCE_LIMIT,
                 memory_count: 0,
@@ -1106,14 +1112,6 @@ impl StoreOpaque {
         &mut self.store_data
     }
 
-    pub fn register_host_trampoline(
-        &mut self,
-        idx: VMSharedSignatureIndex,
-        trampoline: VMTrampoline,
-    ) {
-        self.host_trampolines.insert(idx, trampoline);
-    }
-
     pub fn interrupt_handle(&self) -> Result<InterruptHandle> {
         if self.engine.config().tunables.interruptable {
             Ok(InterruptHandle {
@@ -1166,17 +1164,89 @@ impl StoreOpaque {
         unsafe { wasmtime_runtime::gc(&self.modules, &mut self.externref_activations_table) }
     }
 
-    pub fn lookup_trampoline(&self, anyfunc: &VMCallerCheckedAnyfunc) -> VMTrampoline {
-        // Look up the trampoline with the store's trampolines (from `Func`).
+    /// Looks up the corresponding `VMTrampoline` which can be used to enter
+    /// wasm given an anyfunc function pointer.
+    ///
+    /// This is a somewhat complicated implementation at this time, unfortnately.
+    /// Trampolines are a sort of side-channel of information which is
+    /// specifically juggled by the `wasmtime` crate in a careful fashion. The
+    /// sources for trampolines are:
+    ///
+    /// * Compiled modules - each compiled module has a trampoline for all
+    ///   signatures of functions that escape the module (e.g. exports and
+    ///   `ref.func`-able functions)
+    /// * `Func::new` - host-defined functions with a dynamic signature get an
+    ///   on-the-fly-compiled trampoline (e.g. JIT-compiled as part of the
+    ///   `Func::new` call).
+    /// * `Func::wrap` - host-defined functions where the trampoline is
+    ///   monomorphized in Rust and compiled by LLVM.
+    ///
+    /// The purpose of this function is that given some wasm function pointer we
+    /// need to find the trampoline for it. For compiled wasm modules this is
+    /// pretty easy, the code pointer of the function pointer will point us
+    /// at a wasm module which has a table of trampolines-by-type that we can
+    /// lookup.
+    ///
+    /// If this lookup fails, however, then we're trying to get the trampoline
+    /// for a wasm function pointer defined by the host. The trampoline isn't
+    /// actually stored in the wasm function pointer itself so we need
+    /// side-channels of information. To achieve this a lazy scheme is
+    /// implemented here based on the assumption that most trampoline lookups
+    /// happen for wasm-defined functions, not host-defined functions.
+    ///
+    /// The `Store` already has a list of all functions in
+    /// `self.store_data().funcs`, it's just not indexed in a nice fashion by
+    /// type index or similar. To solve this there's an internal map in each
+    /// store, `host_trampolines`, which maps from a type index to the
+    /// store-owned trampoline. The actual population of this map, however, is
+    /// deferred to this function itself.
+    ///
+    /// Most of the time we are looking up a Wasm function's trampoline when
+    /// calling this function, and we don't want to make insertion of a host
+    /// function into the store more expensive than it has to be. We could
+    /// update the `host_trampolines` whenever a host function is inserted into
+    /// the store, but this is a relatively expensive hash map insertion.
+    /// Instead the work is deferred until we actually look up that trampoline
+    /// in this method.
+    ///
+    /// This all means that if the lookup of the trampoline fails within
+    /// `self.host_trampolines` we lazily populate `self.host_trampolines` by
+    /// iterating over `self.store_data().funcs`, inserting trampolines as we
+    /// go. If we find the right trampoline then it's returned.
+    pub fn lookup_trampoline(&mut self, anyfunc: &VMCallerCheckedAnyfunc) -> VMTrampoline {
+        // First try to see if the `anyfunc` belongs to any module. Each module
+        // has its own map of trampolines-per-type-index and the code pointer in
+        // the `anyfunc` will enable us to quickly find a module.
+        if let Some(trampoline) = self.modules.lookup_trampoline(anyfunc) {
+            return trampoline;
+        }
+
+        // Next consult the list of store-local host trampolines. This is
+        // primarily populated by functions created by `Func::new` or similar
+        // creation functions, host-defined functions.
         if let Some(trampoline) = self.host_trampolines.get(&anyfunc.type_index) {
             return *trampoline;
         }
 
-        // Look up the trampoline with the registered modules
-        if let Some(trampoline) = self.modules.lookup_trampoline(anyfunc) {
-            return trampoline;
+        // If no trampoline was found then it means that it hasn't been loaded
+        // into `host_trampolines` yet. Skip over all the ones we've looked at
+        // so far and start inserting into `self.host_trampolines`, returning
+        // the actual trampoline once found.
+        for f in self
+            .store_data
+            .funcs()
+            .skip(self.host_func_trampolines_registered)
+        {
+            self.host_func_trampolines_registered += 1;
+            self.host_trampolines.insert(f.sig_index(), f.trampoline());
+            if f.sig_index() == anyfunc.type_index {
+                return f.trampoline();
+            }
         }
 
+        // If reached this is a bug in Wasmtime. Lookup of a trampoline should
+        // only happen for wasm functions or host functions, all of which should
+        // be indexed by the above.
         panic!("trampoline missing")
     }
 

crates/wasmtime/src/store/data.rs

@@ -99,6 +99,10 @@ impl StoreData {
         }
         true
     }
+
+    pub(crate) fn funcs(&self) -> impl Iterator<Item = &crate::func::FuncData> {
+        self.funcs.iter()
+    }
 }
 
 impl<T> Index<Stored<T>> for StoreData