Move the Store::signature_cache field (#847)

This commit removes the `signature_cache` field from the `Store` type
and performs a few internal changes which are aimed to be a bit forward
looking towards #777, making `Store` threadsafe.

The changes made here are:

* The `SignatureRegistry` internal type now contains the reverse map
  that `signature_cache` was serving to do. This is populated on calls
  to `register` automatically and is accompanied by a `lookup` method as
  well.

* The `register_wasmtime_signature` and `lookup_wasmtime_signature`
  methods were removed from `Store` and now instead work by using the
  `Compiler::signatures` field.

* The `SignatureRegistry` type was updated to have interior mutability.
  The global `Compiler` type is highly likely to get shared across many
  threads through `Store`, so it needs some form of lock somewhere for
  mutation of the registry of signatures and this commit opts to put it
  inside `SignatureRegistry` which will eventually allow for the removal
  of most `&mut self` method on `Compiler`.

Author: alexcrichton

crates/api/src/instance.rs

@@ -186,7 +186,7 @@ impl Instance {
                 // HACK ensure all handles, instantiated outside Store, present in
                 // the store's SignatureRegistry, e.g. WASI instances that are
                 // imported into this store using the from_handle() method.
-                let _ = store.register_wasmtime_signature(signature);
+                store.compiler().signatures().register(signature);
             }
 
             // We should support everything supported by wasmtime_runtime, or

crates/api/src/module.rs

@@ -401,11 +401,5 @@ fn compile(store: &Store, binary: &[u8], module_name: Option<&str>) -> Result<Co
         module_name,
         store.engine().config().debug_info,
     )?;
-
-    // Register all module signatures
-    for signature in compiled_module.module().signatures.values() {
-        store.register_wasmtime_signature(signature);
-    }
-
     Ok(compiled_module)
 }

crates/api/src/runtime.rs

@@ -1,12 +1,8 @@
 use anyhow::Result;
 use std::cell::RefCell;
-use std::collections::HashMap;
 use std::rc::Rc;
 use std::sync::Arc;
-use wasmtime_environ::{
-    ir,
-    settings::{self, Configurable},
-};
+use wasmtime_environ::settings::{self, Configurable};
 use wasmtime_jit::{native, CompilationStrategy, Compiler, Features};
 
 // Runtime Environment
@@ -342,7 +338,6 @@ pub struct Store {
 struct StoreInner {
     engine: Engine,
     compiler: RefCell<Compiler>,
-    signature_cache: RefCell<HashMap<wasmtime_runtime::VMSharedSignatureIndex, ir::Signature>>,
 }
 
 impl Store {
@@ -354,7 +349,6 @@ impl Store {
             inner: Rc::new(StoreInner {
                 engine: engine.clone(),
                 compiler: RefCell::new(compiler),
-                signature_cache: RefCell::new(HashMap::new()),
             }),
         }
     }
@@ -364,34 +358,12 @@ impl Store {
         &self.inner.engine
     }
 
-    pub(crate) fn compiler_mut(&self) -> std::cell::RefMut<'_, Compiler> {
-        self.inner.compiler.borrow_mut()
-    }
-
-    pub(crate) fn register_wasmtime_signature(
-        &self,
-        signature: &ir::Signature,
-    ) -> wasmtime_runtime::VMSharedSignatureIndex {
-        use std::collections::hash_map::Entry;
-        let index = self.compiler_mut().signatures().register(signature);
-        match self.inner.signature_cache.borrow_mut().entry(index) {
-            Entry::Vacant(v) => {
-                v.insert(signature.clone());
-            }
-            Entry::Occupied(_) => (),
-        }
-        index
+    pub(crate) fn compiler(&self) -> std::cell::Ref<'_, Compiler> {
+        self.inner.compiler.borrow()
     }
 
-    pub(crate) fn lookup_wasmtime_signature(
-        &self,
-        type_index: wasmtime_runtime::VMSharedSignatureIndex,
-    ) -> Option<ir::Signature> {
-        self.inner
-            .signature_cache
-            .borrow()
-            .get(&type_index)
-            .cloned()
+    pub(crate) fn compiler_mut(&self) -> std::cell::RefMut<'_, Compiler> {
+        self.inner.compiler.borrow_mut()
     }
 
     /// Returns whether the stores `a` and `b` refer to the same underlying

crates/api/src/trampoline/create_handle.rs

@@ -12,7 +12,7 @@ use wasmtime_runtime::{Imports, InstanceHandle, VMFunctionBody};
 
 pub(crate) fn create_handle(
     module: Module,
-    signature_registry: Option<&Store>,
+    store: Option<&Store>,
     finished_functions: PrimaryMap<DefinedFuncIndex, *const VMFunctionBody>,
     state: Box<dyn Any>,
 ) -> Result<InstanceHandle> {
@@ -26,12 +26,12 @@ pub(crate) fn create_handle(
     let data_initializers = Vec::new();
 
     // Compute indices into the shared signature table.
-    let signatures = signature_registry
-        .map(|signature_registry| {
+    let signatures = store
+        .map(|store| {
             module
                 .signatures
                 .values()
-                .map(|sig| signature_registry.register_wasmtime_signature(sig))
+                .map(|sig| store.compiler().signatures().register(sig))
                 .collect::<PrimaryMap<_, _>>()
         })
         .unwrap_or_else(PrimaryMap::new);

crates/api/src/values.rs

@@ -205,7 +205,7 @@ pub(crate) fn into_checked_anyfunc(
                 } => (*vmctx, *address, signature),
                 _ => panic!("expected function export"),
             };
-            let type_index = store.register_wasmtime_signature(signature);
+            let type_index = store.compiler().signatures().register(signature);
             wasmtime_runtime::VMCallerCheckedAnyfunc {
                 func_ptr,
                 type_index,
@@ -224,7 +224,9 @@ pub(crate) fn from_checked_anyfunc(
         return Val::AnyRef(AnyRef::Null);
     }
     let signature = store
-        .lookup_wasmtime_signature(item.type_index)
+        .compiler()
+        .signatures()
+        .lookup(item.type_index)
         .expect("signature");
     let instance_handle = unsafe { wasmtime_runtime::InstanceHandle::from_vmctx(item.vmctx) };
     let export = wasmtime_runtime::Export::Function {

crates/jit/src/compiler.rs

@@ -265,8 +265,8 @@ impl Compiler {
     }
 
     /// Shared signature registry.
-    pub fn signatures(&mut self) -> &mut SignatureRegistry {
-        &mut self.signatures
+    pub fn signatures(&self) -> &SignatureRegistry {
+        &self.signatures
     }
 }
 

crates/runtime/src/sig_registry.rs

@@ -5,6 +5,7 @@ use crate::vmcontext::VMSharedSignatureIndex;
 use more_asserts::{assert_lt, debug_assert_lt};
 use std::collections::{hash_map, HashMap};
 use std::convert::TryFrom;
+use std::sync::RwLock;
 use wasmtime_environ::ir;
 
 /// WebAssembly requires that the caller and callee signatures in an indirect
@@ -13,21 +14,33 @@ use wasmtime_environ::ir;
 /// index comparison.
 #[derive(Debug)]
 pub struct SignatureRegistry {
-    signature_hash: HashMap<ir::Signature, VMSharedSignatureIndex>,
+    // This structure is stored in a `Compiler` and is intended to be shared
+    // across many instances. Ideally instances can themselves be sent across
+    // threads, and ideally we can compile across many threads. As a result we
+    // use interior mutability here with a lock to avoid having callers to
+    // externally synchronize calls to compilation.
+    inner: RwLock<Inner>,
+}
+
+#[derive(Debug, Default)]
+struct Inner {
+    signature2index: HashMap<ir::Signature, VMSharedSignatureIndex>,
+    index2signature: HashMap<VMSharedSignatureIndex, ir::Signature>,
 }
 
 impl SignatureRegistry {
     /// Create a new `SignatureRegistry`.
     pub fn new() -> Self {
         Self {
-            signature_hash: HashMap::new(),
+            inner: Default::default(),
         }
     }
 
     /// Register a signature and return its unique index.
-    pub fn register(&mut self, sig: &ir::Signature) -> VMSharedSignatureIndex {
-        let len = self.signature_hash.len();
-        match self.signature_hash.entry(sig.clone()) {
+    pub fn register(&self, sig: &ir::Signature) -> VMSharedSignatureIndex {
+        let mut inner = self.inner.write().unwrap();
+        let len = inner.signature2index.len();
+        match inner.signature2index.entry(sig.clone()) {
             hash_map::Entry::Occupied(entry) => *entry.get(),
             hash_map::Entry::Vacant(entry) => {
                 // Keep `signature_hash` len under 2**32 -- VMSharedSignatureIndex::new(std::u32::MAX)
@@ -39,8 +52,22 @@ impl SignatureRegistry {
                 );
                 let sig_id = VMSharedSignatureIndex::new(u32::try_from(len).unwrap());
                 entry.insert(sig_id);
+                inner.index2signature.insert(sig_id, sig.clone());
                 sig_id
             }
         }
     }
+
+    /// Looks up a shared signature index within this registry.
+    ///
+    /// Note that for this operation to be semantically correct the `idx` must
+    /// have previously come from a call to `register` of this same object.
+    pub fn lookup(&self, idx: VMSharedSignatureIndex) -> Option<ir::Signature> {
+        self.inner
+            .read()
+            .unwrap()
+            .index2signature
+            .get(&idx)
+            .cloned()
+    }
 }