Make garbage collection an async function

crates/wasmtime/src/runtime/func.rs

@@ -2206,11 +2206,11 @@ impl<T> Caller<'_, T> {
     ///
     /// Same as [`Store::gc_async`](crate::Store::gc_async).
     #[cfg(all(feature = "async", feature = "gc"))]
-    pub async fn gc_async(&mut self, why: Option<&crate::GcHeapOutOfMemory<()>>) -> Result<()>
+    pub async fn gc_async(&mut self, why: Option<&crate::GcHeapOutOfMemory<()>>)
     where
         T: Send + 'static,
     {
-        self.store.gc_async(why).await
+        self.store.gc_async(why).await;
     }
 
     /// Returns the remaining fuel in the store.

crates/wasmtime/src/runtime/store.rs

@@ -93,7 +93,8 @@ use crate::runtime::vm::mpk::ProtectionKey;
 use crate::runtime::vm::{
     self, GcStore, Imports, InstanceAllocationRequest, InstanceAllocator, InstanceHandle,
     Interpreter, InterpreterRef, ModuleRuntimeInfo, OnDemandInstanceAllocator, SendSyncPtr,
-    SignalHandler, StoreBox, StorePtr, Unwind, VMContext, VMFuncRef, VMGcRef, VMStoreContext,
+    SignalHandler, StoreBox, StorePtr, Unwind, VMContext, VMFuncRef, VMGcRef, VMStore,
+    VMStoreContext,
 };
 use crate::trampoline::VMHostGlobalContext;
 use crate::{Engine, Module, Trap, Val, ValRaw, module::ModuleRegistry};
@@ -881,8 +882,7 @@ impl<T> Store<T> {
     /// This method is only available when the `gc` Cargo feature is enabled.
     #[cfg(feature = "gc")]
     pub fn gc(&mut self, why: Option<&crate::GcHeapOutOfMemory<()>>) {
-        assert!(!self.inner.async_support());
-        self.inner.gc(why);
+        StoreContextMut(&mut self.inner).gc(why)
     }
 
     /// Returns the amount fuel in this [`Store`]. When fuel is enabled, it must
@@ -1101,7 +1101,9 @@ impl<'a, T> StoreContextMut<'a, T> {
     /// This method is only available when the `gc` Cargo feature is enabled.
     #[cfg(feature = "gc")]
     pub fn gc(&mut self, why: Option<&crate::GcHeapOutOfMemory<()>>) {
-        self.0.gc(why);
+        assert!(!self.0.async_support());
+        let (mut limiter, store) = self.0.resource_limiter_and_store_opaque();
+        vm::assert_ready(store.gc(limiter.as_mut(), None, why.map(|e| e.bytes_needed())));
     }
 
     /// Returns remaining fuel in this store.
@@ -1511,7 +1513,7 @@ impl StoreOpaque {
         #[cfg(feature = "gc")]
         fn allocate_gc_store(
             engine: &Engine,
-            vmstore: NonNull<dyn vm::VMStore>,
+            vmstore: NonNull<dyn VMStore>,
             pkey: Option<ProtectionKey>,
         ) -> Result<GcStore> {
             use wasmtime_environ::packed_option::ReservedValue;
@@ -1558,7 +1560,7 @@ impl StoreOpaque {
         #[cfg(not(feature = "gc"))]
         fn allocate_gc_store(
             _engine: &Engine,
-            _vmstore: NonNull<dyn vm::VMStore>,
+            _vmstore: NonNull<dyn VMStore>,
             _pkey: Option<ProtectionKey>,
         ) -> Result<GcStore> {
             bail!("cannot allocate a GC store: the `gc` feature was disabled at compile time")
@@ -1656,12 +1658,7 @@ impl StoreOpaque {
     }
 
     #[cfg(feature = "gc")]
-    fn do_gc(&mut self) {
-        assert!(
-            !self.async_support(),
-            "must use `store.gc_async()` instead of `store.gc()` for async stores"
-        );
-
+    async fn do_gc(&mut self) {
         // If the GC heap hasn't been initialized, there is nothing to collect.
         if self.gc_store.is_none() {
             return;
@@ -1673,8 +1670,11 @@ impl StoreOpaque {
         // call mutable methods on `self`.
         let mut roots = core::mem::take(&mut self.gc_roots_list);
 
-        self.trace_roots(&mut roots);
-        self.unwrap_gc_store_mut().gc(unsafe { roots.iter() });
+        self.trace_roots(&mut roots).await;
+        let async_yield = self.async_support();
+        self.unwrap_gc_store_mut()
+            .gc(async_yield, unsafe { roots.iter() })
+            .await;
 
         // Restore the GC roots for the next GC.
         roots.clear();
@@ -1684,16 +1684,30 @@ impl StoreOpaque {
     }
 
     #[cfg(feature = "gc")]
-    fn trace_roots(&mut self, gc_roots_list: &mut GcRootsList) {
+    async fn trace_roots(&mut self, gc_roots_list: &mut GcRootsList) {
         log::trace!("Begin trace GC roots");
 
         // We shouldn't have any leftover, stale GC roots.
         assert!(gc_roots_list.is_empty());
 
         self.trace_wasm_stack_roots(gc_roots_list);
+        #[cfg(feature = "async")]
+        if self.async_support() {
+            vm::Yield::new().await;
+        }
         #[cfg(feature = "stack-switching")]
-        self.trace_wasm_continuation_roots(gc_roots_list);
+        {
+            self.trace_wasm_continuation_roots(gc_roots_list);
+            #[cfg(feature = "async")]
+            if self.async_support() {
+                vm::Yield::new().await;
+            }
+        }
         self.trace_vmctx_roots(gc_roots_list);
+        #[cfg(feature = "async")]
+        if self.async_support() {
+            vm::Yield::new().await;
+        }
         self.trace_user_roots(gc_roots_list);
 
         log::trace!("End trace GC roots")
@@ -1927,7 +1941,7 @@ impl StoreOpaque {
     }
 
     #[inline]
-    pub fn traitobj(&self) -> NonNull<dyn vm::VMStore> {
+    pub fn traitobj(&self) -> NonNull<dyn VMStore> {
         self.traitobj.as_raw().unwrap()
     }
 
@@ -2271,7 +2285,7 @@ pub(crate) enum AllocateInstanceKind<'a> {
     },
 }
 
-unsafe impl<T> vm::VMStore for StoreInner<T> {
+unsafe impl<T> VMStore for StoreInner<T> {
     #[cfg(feature = "component-model-async")]
     fn component_async_store(
         &mut self,
@@ -2559,7 +2573,7 @@ impl AsStoreOpaque for StoreOpaque {
     }
 }
 
-impl AsStoreOpaque for dyn vm::VMStore {
+impl AsStoreOpaque for dyn VMStore {
     fn as_store_opaque(&mut self) -> &mut StoreOpaque {
         self
     }

crates/wasmtime/src/runtime/store/async_.rs

@@ -2,6 +2,8 @@
 use crate::CallHook;
 use crate::fiber::{self};
 use crate::prelude::*;
+#[cfg(feature = "gc")]
+use crate::runtime::vm::VMStore;
 use crate::store::{ResourceLimiterInner, StoreInner, StoreOpaque};
 use crate::{Store, StoreContextMut, UpdateDeadline};
 
@@ -89,11 +91,11 @@ impl<T> Store<T> {
     ///
     /// This method is only available when the `gc` Cargo feature is enabled.
     #[cfg(feature = "gc")]
-    pub async fn gc_async(&mut self, why: Option<&crate::GcHeapOutOfMemory<()>>) -> Result<()>
+    pub async fn gc_async(&mut self, why: Option<&crate::GcHeapOutOfMemory<()>>)
     where
         T: Send,
     {
-        self.inner.gc_async(why).await
+        StoreContextMut(&mut self.inner).gc_async(why).await
     }
 
     /// Configures epoch-deadline expiration to yield to the async
@@ -132,11 +134,14 @@ impl<'a, T> StoreContextMut<'a, T> {
     ///
     /// This method is only available when the `gc` Cargo feature is enabled.
     #[cfg(feature = "gc")]
-    pub async fn gc_async(&mut self, why: Option<&crate::GcHeapOutOfMemory<()>>) -> Result<()>
+    pub async fn gc_async(&mut self, why: Option<&crate::GcHeapOutOfMemory<()>>)
     where
         T: Send + 'static,
     {
-        self.0.gc_async(why).await
+        let (mut limiter, store) = self.0.resource_limiter_and_store_opaque();
+        store
+            .gc(limiter.as_mut(), None, why.map(|e| e.bytes_needed()))
+            .await;
     }
 
     /// Configures epoch-deadline expiration to yield to the async
@@ -164,77 +169,6 @@ impl<T> StoreInner<T> {
 
 #[doc(hidden)]
 impl StoreOpaque {
-    /// Executes a synchronous computation `func` asynchronously on a new fiber.
-    ///
-    /// This function will convert the synchronous `func` into an asynchronous
-    /// future. This is done by running `func` in a fiber on a separate native
-    /// stack which can be suspended and resumed from.
-    #[cfg(feature = "gc")]
-    pub(crate) async fn on_fiber<R: Send + Sync>(
-        &mut self,
-        func: impl FnOnce(&mut Self) -> R + Send + Sync,
-    ) -> Result<R> {
-        fiber::on_fiber(self, func).await
-    }
-
-    #[cfg(feature = "gc")]
-    pub(super) async fn do_gc_async(&mut self) {
-        assert!(
-            self.async_support(),
-            "cannot use `gc_async` without enabling async support in the config",
-        );
-
-        // If the GC heap hasn't been initialized, there is nothing to collect.
-        if self.gc_store.is_none() {
-            return;
-        }
-
-        log::trace!("============ Begin Async GC ===========");
-
-        // Take the GC roots out of `self` so we can borrow it mutably but still
-        // call mutable methods on `self`.
-        let mut roots = core::mem::take(&mut self.gc_roots_list);
-
-        self.trace_roots_async(&mut roots).await;
-        self.unwrap_gc_store_mut()
-            .gc_async(unsafe { roots.iter() })
-            .await;
-
-        // Restore the GC roots for the next GC.
-        roots.clear();
-        self.gc_roots_list = roots;
-
-        log::trace!("============ End Async GC ===========");
-    }
-
-    #[inline]
-    #[cfg(not(feature = "gc"))]
-    pub async fn gc_async(&mut self) {
-        // Nothing to collect.
-        //
-        // Note that this is *not* a public method, this is just defined for the
-        // crate-internal `StoreOpaque` type. This is a convenience so that we
-        // don't have to `cfg` every call site.
-    }
-
-    #[cfg(feature = "gc")]
-    async fn trace_roots_async(&mut self, gc_roots_list: &mut crate::runtime::vm::GcRootsList) {
-        use crate::runtime::vm::Yield;
-
-        log::trace!("Begin trace GC roots");
-
-        // We shouldn't have any leftover, stale GC roots.
-        assert!(gc_roots_list.is_empty());
-
-        self.trace_wasm_stack_roots(gc_roots_list);
-        Yield::new().await;
-        self.trace_vmctx_roots(gc_roots_list);
-        Yield::new().await;
-        self.trace_user_roots(gc_roots_list);
-
-        log::trace!("End trace GC roots")
-    }
-
     /// Yields execution to the caller on out-of-gas or epoch interruption.
     ///
     /// This only works on async futures and stores, and assumes that we're