Make table growth a true async fn

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

@@ -1,11 +1,13 @@
 use crate::prelude::*;
 use crate::runtime::RootedGcRefImpl;
-use crate::runtime::vm::{self as runtime, GcStore, TableElementType, VMFuncRef, VMGcRef};
+use crate::runtime::vm::{
+    self, GcStore, SendSyncPtr, TableElementType, VMFuncRef, VMGcRef, VMStore,
+};
 use crate::store::{AutoAssertNoGc, StoreInstanceId, StoreOpaque};
 use crate::trampoline::generate_table_export;
 use crate::{
-    AnyRef, AsContext, AsContextMut, ExnRef, ExternRef, Func, HeapType, Ref, RefType, TableType,
-    Trap,
+    AnyRef, AsContext, AsContextMut, ExnRef, ExternRef, Func, HeapType, Ref, RefType,
+    StoreContextMut, TableType, Trap,
 };
 use core::iter;
 use core::ptr::NonNull;
@@ -139,7 +141,7 @@ impl Table {
         TableType::from_wasmtime_table(store.engine(), self.wasmtime_ty(store))
     }
 
-    /// Returns the `runtime::Table` within `store` as well as the optional
+    /// Returns the `vm::Table` within `store` as well as the optional
     /// `GcStore` in use within `store`.
     ///
     /// # Panics
@@ -149,7 +151,7 @@ impl Table {
         &self,
         store: &'a mut StoreOpaque,
         lazy_init_range: impl IntoIterator<Item = u64>,
-    ) -> (&'a mut runtime::Table, Option<&'a mut GcStore>) {
+    ) -> (&'a mut vm::Table, Option<&'a mut GcStore>) {
         self.instance.assert_belongs_to(store.id());
         let (store, instance) = store.optional_gc_store_and_instance_mut(self.instance.instance());
 
@@ -281,42 +283,60 @@ impl Table {
     /// When using an async resource limiter, use [`Table::grow_async`]
     /// instead.
     pub fn grow(&self, mut store: impl AsContextMut, delta: u64, init: Ref) -> Result<u64> {
-        let store = store.as_context_mut().0;
+        vm::one_poll(self._grow(store.as_context_mut(), delta, init))
+            .expect("must use `grow_async` when async resource limiters are in use")
+    }
+
+    async fn _grow<T>(&self, store: StoreContextMut<'_, T>, delta: u64, init: Ref) -> Result<u64> {
+        let store = store.0;
         let ty = self.ty(&store);
-        let (table, _gc_store) = self.wasmtime_table(store, iter::empty());
-        // FIXME(#11179) shouldn't need to subvert the borrow checker
-        let table: *mut _ = table;
-        unsafe {
-            let result = match element_type(&ty) {
-                TableElementType::Func => {
-                    let element = init.into_table_func(store, ty.element())?;
-                    (*table).grow_func(store, delta, element)?
-                }
-                TableElementType::GcRef => {
-                    // FIXME: `grow_gc_ref` shouldn't require the whole store
-                    // and should require `AutoAssertNoGc`. For now though we
-                    // know that table growth doesn't trigger GC so it should be
-                    // ok to create a copy of the GC reference even though it's
-                    // not tracked anywhere.
-                    let element = init
-                        .into_table_gc_ref(&mut AutoAssertNoGc::new(store), ty.element())?
-                        .map(|r| r.unchecked_copy());
-                    (*table).grow_gc_ref(store, delta, element.as_ref())?
-                }
-                // TODO(#10248) Required to support stack switching in the
-                // embedder API.
-                TableElementType::Cont => bail!("unimplemented table for cont"),
-            };
-            match result {
-                Some(size) => {
-                    let vm = (*table).vmtable();
-                    store[self.instance].table_ptr(self.index).write(vm);
-                    // unwrap here should be ok because the runtime should always guarantee
-                    // that we can fit the table size in a 64-bit integer.
-                    Ok(u64::try_from(size).unwrap())
-                }
-                None => bail!("failed to grow table by `{}`", delta),
+        let (mut limiter, store) = store.resource_limiter_and_store_opaque();
+        let limiter = limiter.as_mut();
+        let result = match element_type(&ty) {
+            TableElementType::Func => {
+                let element = init
+                    .into_table_func(store, ty.element())?
+                    .map(SendSyncPtr::new);
+                self.instance
+                    .get_mut(store)
+                    .defined_table_grow(self.index, async |table| {
+                        // SAFETY: in the context of `defined_table_grow` this
+                        // is safe to call as it'll update the internal table
+                        // pointer in the instance.
+                        unsafe { table.grow_func(limiter, delta, element).await }
+                    })
+                    .await?
             }
+            TableElementType::GcRef => {
+                let mut store = AutoAssertNoGc::new(store);
+                let element = init
+                    .into_table_gc_ref(&mut store, ty.element())?
+                    .map(|r| r.unchecked_copy());
+                let (gc_store, instance) = self.instance.get_with_gc_store_mut(&mut store);
+                instance
+                    .defined_table_grow(self.index, async |table| {
+                        // SAFETY: in the context of `defined_table_grow` this
+                        // is safe to call as it'll update the internal table
+                        // pointer in the instance.
+                        unsafe {
+                            table
+                                .grow_gc_ref(limiter, gc_store, delta, element.as_ref())
+                                .await
+                        }
+                    })
+                    .await?
+            }
+            // TODO(#10248) Required to support stack switching in the
+            // embedder API.
+            TableElementType::Cont => bail!("unimplemented table for cont"),
+        };
+        match result {
+            Some(size) => {
+                // unwrap here should be ok because the runtime should always
+                // guarantee that we can fit the table size in a 64-bit integer.
+                Ok(u64::try_from(size).unwrap())
+            }
+            None => bail!("failed to grow table by `{}`", delta),
         }
     }
 
@@ -334,14 +354,7 @@ impl Table {
         delta: u64,
         init: Ref,
     ) -> Result<u64> {
-        let mut store = store.as_context_mut();
-        assert!(
-            store.0.async_support(),
-            "cannot use `grow_async` without enabling async support on the config"
-        );
-        store
-            .on_fiber(|store| self.grow(store, delta, init))
-            .await?
+        self._grow(store.as_context_mut(), delta, init).await
     }
 
     /// Copy `len` elements from `src_table[src_index..]` into
@@ -516,11 +529,7 @@ impl Table {
     }
 
     #[cfg(feature = "gc")]
-    pub(crate) fn trace_roots(
-        &self,
-        store: &mut StoreOpaque,
-        gc_roots_list: &mut crate::runtime::vm::GcRootsList,
-    ) {
+    pub(crate) fn trace_roots(&self, store: &mut StoreOpaque, gc_roots_list: &mut vm::GcRootsList) {
         if !self
             ._ty(store)
             .element()
@@ -549,9 +558,9 @@ impl Table {
         &module.tables[index]
     }
 
-    pub(crate) fn vmimport(&self, store: &StoreOpaque) -> crate::runtime::vm::VMTableImport {
+    pub(crate) fn vmimport(&self, store: &StoreOpaque) -> vm::VMTableImport {
         let instance = &store[self.instance];
-        crate::runtime::vm::VMTableImport {
+        vm::VMTableImport {
             from: instance.table_ptr(self.index).into(),
             vmctx: instance.vmctx().into(),
             index: self.index,
@@ -691,4 +700,12 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn grow_is_send() {
+        fn _assert_send<T: Send>(_: T) {}
+        fn _grow(table: &Table, store: &mut Store<()>, init: Ref) {
+            _assert_send(table.grow(store, 0, init))
+        }
+    }
 }

crates/wasmtime/src/runtime/fiber.rs

@@ -393,7 +393,6 @@ impl StoreOpaque {
     /// # Panics
     ///
     /// Panics if this is invoked outside the context of a fiber.
-    #[cfg(feature = "component-model-async")]
     pub(crate) fn with_blocking<R>(
         &mut self,
         f: impl FnOnce(&mut Self, &mut BlockingContext<'_, '_>) -> R,

crates/wasmtime/src/runtime/limits.rs

@@ -29,7 +29,7 @@ pub const DEFAULT_MEMORY_LIMIT: usize = 10000;
 /// or not and you're otherwise working in an asynchronous context the
 /// [`ResourceLimiterAsync`] trait is also provided to avoid blocking an OS
 /// thread while a limit is determined.
-pub trait ResourceLimiter {
+pub trait ResourceLimiter: Send {
     /// Notifies the resource limiter that an instance's linear memory has been
     /// requested to grow.
     ///
@@ -172,7 +172,7 @@ pub trait ResourceLimiter {
 /// answer the question whether growing a memory or table is allowed.
 #[cfg(feature = "async")]
 #[async_trait::async_trait]
-pub trait ResourceLimiterAsync {
+pub trait ResourceLimiterAsync: Send {
     /// Async version of [`ResourceLimiter::memory_growing`]
     async fn memory_growing(
         &mut self,

crates/wasmtime/src/runtime/store.rs

@@ -247,6 +247,43 @@ enum ResourceLimiterInner<T> {
     Async(Box<dyn (FnMut(&mut T) -> &mut dyn crate::ResourceLimiterAsync) + Send + Sync>),
 }
 
+/// Representation of a configured resource limiter for a store.
+///
+/// This is acquired with `resource_limiter_and_store_opaque` for example and is
+/// threaded through to growth operations on tables/memories. Note that this is
+/// passed around as `Option<&mut StoreResourceLimiter<'_>>` to make it
+/// efficient to pass around (nullable pointer) and it's also notably passed
+/// around as an `Option` to represent how this is optionally specified within a
+/// store.
+pub enum StoreResourceLimiter<'a> {
+    Sync(&'a mut dyn crate::ResourceLimiter),
+    #[cfg(feature = "async")]
+    Async(&'a mut dyn crate::ResourceLimiterAsync),
+}
+
+impl StoreResourceLimiter<'_> {
+    pub(crate) async fn table_growing(
+        &mut self,
+        current: usize,
+        desired: usize,
+        maximum: Option<usize>,
+    ) -> Result<bool, Error> {
+        match self {
+            Self::Sync(s) => s.table_growing(current, desired, maximum),
+            #[cfg(feature = "async")]
+            Self::Async(s) => s.table_growing(current, desired, maximum).await,
+        }
+    }
+
+    pub(crate) fn table_grow_failed(&mut self, error: anyhow::Error) -> Result<()> {
+        match self {
+            Self::Sync(s) => s.table_grow_failed(error),
+            #[cfg(feature = "async")]
+            Self::Async(s) => s.table_grow_failed(error),
+        }
+    }
+}
+
 enum CallHookInner<T: 'static> {
     #[cfg(feature = "call-hook")]
     Sync(Box<dyn FnMut(StoreContextMut<'_, T>, CallHook) -> Result<()> + Send + Sync>),
@@ -2234,6 +2271,19 @@ unsafe impl<T> vm::VMStore for StoreInner<T> {
         &mut self.inner
     }
 
+    fn resource_limiter_and_store_opaque(
+        &mut self,
+    ) -> (Option<StoreResourceLimiter<'_>>, &mut StoreOpaque) {
+        (
+            self.limiter.as_mut().map(|l| match l {
+                ResourceLimiterInner::Sync(s) => StoreResourceLimiter::Sync(s(&mut self.data)),
+                #[cfg(feature = "async")]
+                ResourceLimiterInner::Async(s) => StoreResourceLimiter::Async(s(&mut self.data)),
+            }),
+            &mut self.inner,
+        )
+    }
+
     fn memory_growing(
         &mut self,
         current: usize,

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

@@ -1,4 +1,4 @@
-use crate::runtime::vm::{self, VMStore};
+use crate::runtime::vm::{self, GcStore, VMStore};
 use crate::store::StoreOpaque;
 use crate::{StoreContext, StoreContextMut};
 use core::num::NonZeroU64;
@@ -252,6 +252,16 @@ impl StoreInstanceId {
         self.assert_belongs_to(store.id());
         store.instance_mut(self.instance)
     }
+
+    /// Same as [`Self::get_mut`], but also returns the `GcStore`.
+    #[inline]
+    pub(crate) fn get_with_gc_store_mut<'a>(
+        &self,
+        store: &'a mut StoreOpaque,
+    ) -> (Option<&'a mut GcStore>, Pin<&'a mut vm::Instance>) {
+        self.assert_belongs_to(store.id());
+        store.optional_gc_store_and_instance_mut(self.instance)
+    }
 }
 
 impl Index<StoreInstanceId> for StoreOpaque {

crates/wasmtime/src/runtime/vm.rs

@@ -35,15 +35,15 @@ pub(crate) struct f64x2(crate::uninhabited::Uninhabited);
 
 use crate::StoreContextMut;
 use crate::prelude::*;
-use crate::store::StoreInner;
-use crate::store::StoreOpaque;
+use crate::store::{StoreInner, StoreOpaque, StoreResourceLimiter};
 use crate::type_registry::RegisteredType;
 use alloc::sync::Arc;
 use core::fmt;
-use core::ops::Deref;
-use core::ops::DerefMut;
+use core::ops::{Deref, DerefMut};
+use core::pin::pin;
 use core::ptr::NonNull;
 use core::sync::atomic::{AtomicUsize, Ordering};
+use core::task::{Context, Poll, Waker};
 use wasmtime_environ::{
     DefinedFuncIndex, DefinedMemoryIndex, HostPtr, VMOffsets, VMSharedTypeIndex,
 };
@@ -201,6 +201,12 @@ pub unsafe trait VMStore: 'static {
     /// Get an exclusive borrow of this store's `StoreOpaque`.
     fn store_opaque_mut(&mut self) -> &mut StoreOpaque;
 
+    /// Returns a split borrow to the limiter plus `StoreOpaque` at the same
+    /// time.
+    fn resource_limiter_and_store_opaque(
+        &mut self,
+    ) -> (Option<StoreResourceLimiter<'_>>, &mut StoreOpaque);
+
     /// Callback invoked to allow the store's resource limiter to reject a
     /// memory grow operation.
     fn memory_growing(
@@ -509,3 +515,38 @@ impl fmt::Display for WasmFault {
         )
     }
 }
+
+/// Asserts that the future `f` is ready and returns its output.
+///
+/// This function is intended to be used when `async_support` is verified as
+/// disabled. Internals of Wasmtime are generally `async` when they optionally
+/// can be, meaning that synchronous entrypoints will invoke this function
+/// after invoking the asynchronous internals. Due to `async_support` being
+/// disabled there should be no way to introduce a yield point meaning that all
+/// futures built from internal functions should always be ready.
+///
+/// # Panics
+///
+/// Panics if `f` is not yet ready.
+pub fn assert_ready<F: Future>(f: F) -> F::Output {
+    one_poll(f).unwrap()
+}
+
+/// Attempts one poll of `f` to see if its output is available.
+///
+/// This function is intended for a few minor entrypoints into the Wasmtime API
+/// where a synchronous function is documented to work even when `async_support`
+/// is enabled. For example growing a `Memory` can be done with a synchronous
+/// function, but it's documented to panic with an async resource limiter.
+///
+/// This function provides the opportunity to poll `f` once to see if its output
+/// is available. If it isn't then `None` is returned and an appropriate panic
+/// message should be generated recommending to use an async function (e.g.
+/// `grow_async` instead of `grow`).
+pub fn one_poll<F: Future>(f: F) -> Option<F::Output> {
+    let mut context = Context::from_waker(&Waker::noop());
+    match pin!(f).poll(&mut context) {
+        Poll::Ready(output) => Some(output),
+        Poll::Pending => None,
+    }
+}

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

@@ -766,13 +766,13 @@ impl Instance {
     /// Performs a grow operation on the `table_index` specified using `grow`.
     ///
     /// This will handle updating the VMTableDefinition internally as necessary.
-    pub(crate) fn defined_table_grow(
+    pub(crate) async fn defined_table_grow(
         mut self: Pin<&mut Self>,
         table_index: DefinedTableIndex,
-        grow: impl FnOnce(&mut Table) -> Result<Option<usize>>,
+        grow: impl AsyncFnOnce(&mut Table) -> Result<Option<usize>>,
     ) -> Result<Option<usize>> {
         let table = self.as_mut().get_defined_table(table_index);
-        let result = grow(table);
+        let result = grow(table).await;
         let element = table.vmtable();
         self.set_table(table_index, element);
         result