Minimize lazy allocation of the GC store (#11411)

* Minimize lazy allocation of the GC store

This commit is an effort to minimize the number of entrypoints which
might lazily allocate a GC store. The is currently done through
`StoreOpaque::gc_store_mut` but this method is very commonly used
meaning that there are many many places to audit for lazily allocating a
GC store. The reason that this needs an audit is that lazy allocation
is an async operation right now that must be on a fiber and is something
I'm looking to fix as part of #11262.

This commit performs a few refactorings to achieve this:

* `gc_store_mut` is renamed to `ensure_gc_store`. This is intended to be
  an `async` function in the future and clearly demarcates where lazy
  allocation of a GC store is occurring.

* `require_gc_store{,_mut}` is now added which is a pure accessor of the
  GC store with no lazy allocation. Most locations previously using
  `gc_store_mut` are updated to use this instead.

Documentation is added to store methods to clearly indicate which ones
are allocating and which ones should only be called in a context where
allocation should already have happened.

* Fix configured build

* Relax GC store restrictions in more places

* Review comments on documentation

* Move `ensure_gc_store` calls during instantiation

Instead update `needs_gc_heap` with the tables that are added to a
module and rely on instantiation to create the GC heap.

* Shuffle around some code

* Fix CI and review comments

* Add in a few more i31 cases for externref

Author: alexcrichton

crates/environ/src/compile/module_environ.rs

@@ -360,6 +360,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                 for entry in tables {
                     let wasmparser::Table { ty, init } = entry?;
                     let table = self.convert_table_type(&ty)?;
+                    self.result.module.needs_gc_heap |= table.ref_type.is_vmgcref_type();
                     self.result.module.tables.push(table);
                     let init = match init {
                         wasmparser::TableInit::RefNull => TableInitialValue::Null {

crates/wasmtime/src/runtime/externals/global.rs

@@ -158,7 +158,7 @@ impl Global {
                         HeapType::NoFunc => Ref::Func(None),
 
                         HeapType::Extern => Ref::Extern(definition.as_gc_ref().map(|r| {
-                            let r = store.unwrap_gc_store_mut().clone_gc_ref(r);
+                            let r = store.clone_gc_ref(r);
                             ExternRef::from_cloned_gc_ref(&mut store, r)
                         })),
 
@@ -180,7 +180,7 @@ impl Global {
                         | HeapType::ConcreteExn(_) => definition
                             .as_gc_ref()
                             .map(|r| {
-                                let r = store.unwrap_gc_store_mut().clone_gc_ref(r);
+                                let r = store.clone_gc_ref(r);
                                 AnyRef::from_cloned_gc_ref(&mut store, r)
                             })
                             .into(),
@@ -236,23 +236,23 @@ impl Global {
                         Some(e) => Some(e.try_gc_ref(&store)?.unchecked_copy()),
                     };
                     let new = new.as_ref();
-                    definition.write_gc_ref(store.unwrap_gc_store_mut(), new);
+                    definition.write_gc_ref(&mut store, new);
                 }
                 Val::AnyRef(a) => {
                     let new = match a {
                         None => None,
                         Some(a) => Some(a.try_gc_ref(&store)?.unchecked_copy()),
                     };
                     let new = new.as_ref();
-                    definition.write_gc_ref(store.unwrap_gc_store_mut(), new);
+                    definition.write_gc_ref(&mut store, new);
                 }
                 Val::ExnRef(e) => {
                     let new = match e {
                         None => None,
                         Some(e) => Some(e.try_gc_ref(&store)?.unchecked_copy()),
                     };
                     let new = new.as_ref();
-                    definition.write_gc_ref(store.unwrap_gc_store_mut(), new);
+                    definition.write_gc_ref(&mut store, new);
                 }
             }
         }

crates/wasmtime/src/runtime/gc/enabled/anyref.rs

@@ -280,11 +280,7 @@ impl AnyRef {
     // (Not actually memory unsafe since we have indexed GC heaps.)
     pub(crate) fn _from_raw(store: &mut AutoAssertNoGc, raw: u32) -> Option<Rooted<Self>> {
         let gc_ref = VMGcRef::from_raw_u32(raw)?;
-        let gc_ref = if gc_ref.is_i31() {
-            gc_ref.copy_i31()
-        } else {
-            store.unwrap_gc_store_mut().clone_gc_ref(&gc_ref)
-        };
+        let gc_ref = store.clone_gc_ref(&gc_ref);
         Some(Self::from_cloned_gc_ref(store, gc_ref))
     }
 
@@ -340,7 +336,7 @@ impl AnyRef {
         let raw = if gc_ref.is_i31() {
             gc_ref.as_raw_non_zero_u32()
         } else {
-            store.gc_store_mut()?.expose_gc_ref_to_wasm(gc_ref)
+            store.require_gc_store_mut()?.expose_gc_ref_to_wasm(gc_ref)
         };
         Ok(raw.get())
     }
@@ -364,7 +360,7 @@ impl AnyRef {
             return Ok(HeapType::I31);
         }
 
-        let header = store.gc_store()?.header(gc_ref);
+        let header = store.require_gc_store()?.header(gc_ref);
 
         if header.kind().matches(VMGcKind::ExternRef) {
             return Ok(HeapType::Any);
@@ -437,7 +433,11 @@ impl AnyRef {
     pub(crate) fn _is_eqref(&self, store: &StoreOpaque) -> Result<bool> {
         assert!(self.comes_from_same_store(store));
         let gc_ref = self.inner.try_gc_ref(store)?;
-        Ok(gc_ref.is_i31() || store.gc_store()?.kind(gc_ref).matches(VMGcKind::EqRef))
+        Ok(gc_ref.is_i31()
+            || store
+                .require_gc_store()?
+                .kind(gc_ref)
+                .matches(VMGcKind::EqRef))
     }
 
     /// Downcast this `anyref` to an `eqref`.
@@ -558,7 +558,11 @@ impl AnyRef {
 
     pub(crate) fn _is_struct(&self, store: &StoreOpaque) -> Result<bool> {
         let gc_ref = self.inner.try_gc_ref(store)?;
-        Ok(!gc_ref.is_i31() && store.gc_store()?.kind(gc_ref).matches(VMGcKind::StructRef))
+        Ok(!gc_ref.is_i31()
+            && store
+                .require_gc_store()?
+                .kind(gc_ref)
+                .matches(VMGcKind::StructRef))
     }
 
     /// Downcast this `anyref` to a `structref`.
@@ -622,7 +626,11 @@ impl AnyRef {
 
     pub(crate) fn _is_array(&self, store: &StoreOpaque) -> Result<bool> {
         let gc_ref = self.inner.try_gc_ref(store)?;
-        Ok(!gc_ref.is_i31() && store.gc_store()?.kind(gc_ref).matches(VMGcKind::ArrayRef))
+        Ok(!gc_ref.is_i31()
+            && store
+                .require_gc_store()?
+                .kind(gc_ref)
+                .matches(VMGcKind::ArrayRef))
     }
 
     /// Downcast this `anyref` to an `arrayref`.

crates/wasmtime/src/runtime/gc/enabled/arrayref.rs

@@ -411,7 +411,7 @@ impl ArrayRef {
         // Allocate the array and write each field value into the appropriate
         // offset.
         let arrayref = store
-            .gc_store_mut()?
+            .require_gc_store_mut()?
             .alloc_uninit_array(allocator.type_index(), len, allocator.layout())
             .context("unrecoverable error when allocating new `arrayref`")?
             .map_err(|n| GcHeapOutOfMemory::new((), n))?;
@@ -432,7 +432,7 @@ impl ArrayRef {
         })() {
             Ok(()) => Ok(Rooted::new(&mut store, arrayref.into())),
             Err(e) => {
-                store.gc_store_mut()?.dealloc_uninit_array(arrayref);
+                store.require_gc_store_mut()?.dealloc_uninit_array(arrayref);
                 Err(e)
             }
         }
@@ -633,7 +633,7 @@ impl ArrayRef {
         assert!(self.comes_from_same_store(store));
         let gc_ref = self.inner.try_gc_ref(store)?;
         debug_assert!({
-            let header = store.gc_store()?.header(gc_ref);
+            let header = store.require_gc_store()?.header(gc_ref);
             header.kind().matches(VMGcKind::ArrayRef)
         });
         let arrayref = gc_ref.as_arrayref_unchecked();
@@ -667,7 +667,7 @@ impl ArrayRef {
         let store = AutoAssertNoGc::new(store);
 
         let gc_ref = self.inner.try_gc_ref(&store)?;
-        let header = store.gc_store()?.header(gc_ref);
+        let header = store.require_gc_store()?.header(gc_ref);
         debug_assert!(header.kind().matches(VMGcKind::ArrayRef));
 
         let len = self._len(&store)?;
@@ -720,7 +720,7 @@ impl ArrayRef {
     fn header<'a>(&self, store: &'a AutoAssertNoGc<'_>) -> Result<&'a VMGcHeader> {
         assert!(self.comes_from_same_store(&store));
         let gc_ref = self.inner.try_gc_ref(store)?;
-        Ok(store.gc_store()?.header(gc_ref))
+        Ok(store.require_gc_store()?.header(gc_ref))
     }
 
     fn arrayref<'a>(&self, store: &'a AutoAssertNoGc<'_>) -> Result<&'a VMArrayRef> {
@@ -843,7 +843,7 @@ impl ArrayRef {
 
     pub(crate) fn type_index(&self, store: &StoreOpaque) -> Result<VMSharedTypeIndex> {
         let gc_ref = self.inner.try_gc_ref(store)?;
-        let header = store.gc_store()?.header(gc_ref);
+        let header = store.require_gc_store()?.header(gc_ref);
         debug_assert!(header.kind().matches(VMGcKind::ArrayRef));
         Ok(header.ty().expect("arrayrefs should have concrete types"))
     }

crates/wasmtime/src/runtime/gc/enabled/eqref.rs

@@ -198,7 +198,7 @@ impl EqRef {
             return Ok(HeapType::I31);
         }
 
-        let header = store.gc_store()?.header(gc_ref);
+        let header = store.require_gc_store()?.header(gc_ref);
 
         if header.kind().matches(VMGcKind::StructRef) {
             return Ok(HeapType::ConcreteStruct(
@@ -320,7 +320,11 @@ impl EqRef {
 
     pub(crate) fn _is_struct(&self, store: &StoreOpaque) -> Result<bool> {
         let gc_ref = self.inner.try_gc_ref(store)?;
-        Ok(!gc_ref.is_i31() && store.gc_store()?.kind(gc_ref).matches(VMGcKind::StructRef))
+        Ok(!gc_ref.is_i31()
+            && store
+                .require_gc_store()?
+                .kind(gc_ref)
+                .matches(VMGcKind::StructRef))
     }
 
     /// Downcast this `eqref` to a `structref`.
@@ -384,7 +388,11 @@ impl EqRef {
 
     pub(crate) fn _is_array(&self, store: &StoreOpaque) -> Result<bool> {
         let gc_ref = self.inner.try_gc_ref(store)?;
-        Ok(!gc_ref.is_i31() && store.gc_store()?.kind(gc_ref).matches(VMGcKind::ArrayRef))
+        Ok(!gc_ref.is_i31()
+            && store
+                .require_gc_store()?
+                .kind(gc_ref)
+                .matches(VMGcKind::ArrayRef))
     }
 
     /// Downcast this `eqref` to an `arrayref`.

crates/wasmtime/src/runtime/gc/enabled/exnref.rs

@@ -166,7 +166,7 @@ impl ExnRef {
     // (Not actually memory unsafe since we have indexed GC heaps.)
     pub(crate) fn _from_raw(store: &mut AutoAssertNoGc, raw: u32) -> Option<Rooted<Self>> {
         let gc_ref = VMGcRef::from_raw_u32(raw)?;
-        let gc_ref = store.unwrap_gc_store_mut().clone_gc_ref(&gc_ref);
+        let gc_ref = store.clone_gc_ref(&gc_ref);
         Some(Self::from_cloned_gc_ref(store, gc_ref))
     }
 
@@ -335,7 +335,7 @@ impl ExnRef {
         // Allocate the exn and write each field value into the appropriate
         // offset.
         let exnref = store
-            .gc_store_mut()?
+            .require_gc_store_mut()?
             .alloc_uninit_exn(allocator.type_index(), &allocator.layout())
             .context("unrecoverable error when allocating new `exnref`")?
             .map_err(|n| GcHeapOutOfMemory::new((), n))?;
@@ -361,15 +361,15 @@ impl ExnRef {
         })() {
             Ok(()) => Ok(Rooted::new(&mut store, exnref.into())),
             Err(e) => {
-                store.gc_store_mut()?.dealloc_uninit_exn(exnref);
+                store.require_gc_store_mut()?.dealloc_uninit_exn(exnref);
                 Err(e)
             }
         }
     }
 
     pub(crate) fn type_index(&self, store: &StoreOpaque) -> Result<VMSharedTypeIndex> {
         let gc_ref = self.inner.try_gc_ref(store)?;
-        let header = store.gc_store()?.header(gc_ref);
+        let header = store.require_gc_store()?.header(gc_ref);
         debug_assert!(header.kind().matches(VMGcKind::ExnRef));
         Ok(header.ty().expect("exnrefs should have concrete types"))
     }
@@ -421,7 +421,7 @@ impl ExnRef {
         let raw = if gc_ref.is_i31() {
             gc_ref.as_raw_non_zero_u32()
         } else {
-            store.gc_store_mut()?.expose_gc_ref_to_wasm(gc_ref)
+            store.require_gc_store_mut()?.expose_gc_ref_to_wasm(gc_ref)
         };
         Ok(raw.get())
     }
@@ -501,7 +501,7 @@ impl ExnRef {
         let store = AutoAssertNoGc::new(store);
 
         let gc_ref = self.inner.try_gc_ref(&store)?;
-        let header = store.gc_store()?.header(gc_ref);
+        let header = store.require_gc_store()?.header(gc_ref);
         debug_assert!(header.kind().matches(VMGcKind::ExnRef));
 
         let index = header.ty().expect("exnrefs should have concrete types");
@@ -554,7 +554,7 @@ impl ExnRef {
     fn header<'a>(&self, store: &'a AutoAssertNoGc<'_>) -> Result<&'a VMGcHeader> {
         assert!(self.comes_from_same_store(&store));
         let gc_ref = self.inner.try_gc_ref(store)?;
-        Ok(store.gc_store()?.header(gc_ref))
+        Ok(store.require_gc_store()?.header(gc_ref))
     }
 
     fn exnref<'a>(&self, store: &'a AutoAssertNoGc<'_>) -> Result<&'a VMExnRef> {

crates/wasmtime/src/runtime/gc/enabled/externref.rs

@@ -229,7 +229,7 @@ impl ExternRef {
         let gc_ref = store
             .retry_after_gc(value, |store, value| {
                 store
-                    .gc_store_mut()?
+                    .require_gc_store_mut()?
                     .alloc_externref(value)
                     .context("unrecoverable error when allocating new `externref`")?
                     .map_err(|(x, n)| GcHeapOutOfMemory::new(x, n).into())
@@ -346,7 +346,7 @@ impl ExternRef {
         let gc_ref = store
             .retry_after_gc_async(value, |store, value| {
                 store
-                    .gc_store_mut()?
+                    .require_gc_store_mut()?
                     .alloc_externref(value)
                     .context("unrecoverable error when allocating new `externref`")?
                     .map_err(|(x, n)| GcHeapOutOfMemory::new(x, n).into())
@@ -435,11 +435,13 @@ impl ExternRef {
         store: &mut AutoAssertNoGc<'_>,
         gc_ref: VMGcRef,
     ) -> Rooted<Self> {
-        assert!(
-            gc_ref.is_extern_ref(&*store.unwrap_gc_store().gc_heap)
-                || gc_ref.is_any_ref(&*store.unwrap_gc_store().gc_heap),
-            "GC reference {gc_ref:#p} should be an externref or anyref"
-        );
+        if !gc_ref.is_i31() {
+            assert!(
+                gc_ref.is_extern_ref(&*store.unwrap_gc_store().gc_heap)
+                    || gc_ref.is_any_ref(&*store.unwrap_gc_store().gc_heap),
+                "GC reference {gc_ref:#p} should be an externref or anyref"
+            );
+        }
         Rooted::new(store, gc_ref)
     }
 
@@ -481,7 +483,10 @@ impl ExternRef {
     {
         let store = store.into().0;
         let gc_ref = self.inner.try_gc_ref(&store)?;
-        let gc_store = store.gc_store()?;
+        if gc_ref.is_i31() {
+            return Ok(None);
+        }
+        let gc_store = store.require_gc_store()?;
         if let Some(externref) = gc_ref.as_externref(&*gc_store.gc_heap) {
             Ok(Some(gc_store.externref_host_data(externref)))
         } else {
@@ -532,7 +537,10 @@ impl ExternRef {
         // so that we can get the store's GC store. But importantly we cannot
         // trigger a GC while we are working with `gc_ref` here.
         let gc_ref = self.inner.try_gc_ref(store)?.unchecked_copy();
-        let gc_store = store.gc_store_mut()?;
+        if gc_ref.is_i31() {
+            return Ok(None);
+        }
+        let gc_store = store.require_gc_store_mut()?;
         if let Some(externref) = gc_ref.as_externref(&*gc_store.gc_heap) {
             Ok(Some(gc_store.externref_host_data_mut(externref)))
         } else {
@@ -579,7 +587,7 @@ impl ExternRef {
     // (Not actually memory unsafe since we have indexed GC heaps.)
     pub(crate) fn _from_raw(store: &mut AutoAssertNoGc, raw: u32) -> Option<Rooted<ExternRef>> {
         let gc_ref = VMGcRef::from_raw_u32(raw)?;
-        let gc_ref = store.unwrap_gc_store_mut().clone_gc_ref(&gc_ref);
+        let gc_ref = store.clone_gc_ref(&gc_ref);
         Some(Self::from_cloned_gc_ref(store, gc_ref))
     }