Make memory growth an async function

crates/wasmtime/src/runtime/memory.rs

@@ -1,6 +1,7 @@
 use crate::Trap;
 use crate::prelude::*;
-use crate::store::{StoreInstanceId, StoreOpaque};
+use crate::runtime::vm::{self, VMStore};
+use crate::store::{StoreInstanceId, StoreOpaque, StoreResourceLimiter};
 use crate::trampoline::generate_memory_export;
 use crate::{AsContext, AsContextMut, Engine, MemoryType, StoreContext, StoreContextMut};
 use core::cell::UnsafeCell;
@@ -596,20 +597,9 @@ impl Memory {
     /// ```
     pub fn grow(&self, mut store: impl AsContextMut, delta: u64) -> Result<u64> {
         let store = store.as_context_mut().0;
-        // FIXME(#11179) shouldn't use a raw pointer to work around the borrow
-        // checker here.
-        let mem: *mut _ = self.wasmtime_memory(store);
-        unsafe {
-            match (*mem).grow(delta, Some(store))? {
-                Some(size) => {
-                    let vm = (*mem).vmmemory();
-                    store[self.instance].memory_ptr(self.index).write(vm);
-                    let page_size = (*mem).page_size();
-                    Ok(u64::try_from(size).unwrap() / page_size)
-                }
-                None => bail!("failed to grow memory by `{}`", delta),
-            }
-        }
+        let (mut limiter, store) = store.resource_limiter_and_store_opaque();
+        vm::one_poll(self._grow(store, limiter.as_mut(), delta))
+            .expect("must use `grow_async` if an async resource limiter is used")
     }
 
     /// Async variant of [`Memory::grow`]. Required when using a
@@ -621,21 +611,29 @@ impl Memory {
     /// [`Store`](`crate::Store`).
     #[cfg(feature = "async")]
     pub async fn grow_async(&self, mut store: impl AsContextMut, delta: u64) -> 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)).await?
+        let store = store.as_context_mut();
+        let (mut limiter, store) = store.0.resource_limiter_and_store_opaque();
+        self._grow(store, limiter.as_mut(), delta).await
     }
 
-    fn wasmtime_memory<'a>(
+    async fn _grow(
         &self,
-        store: &'a mut StoreOpaque,
-    ) -> &'a mut crate::runtime::vm::Memory {
-        self.instance
+        store: &mut StoreOpaque,
+        limiter: Option<&mut StoreResourceLimiter<'_>>,
+        delta: u64,
+    ) -> Result<u64> {
+        let result = self
+            .instance
             .get_mut(store)
-            .get_defined_memory_mut(self.index)
+            .memory_grow(limiter, self.index, delta)
+            .await?;
+        match result {
+            Some(size) => {
+                let page_size = self.wasmtime_ty(store).page_size();
+                Ok(u64::try_from(size).unwrap() / page_size)
+            }
+            None => bail!("failed to grow memory by `{}`", delta),
+        }
     }
 
     pub(crate) fn from_raw(instance: StoreInstanceId, index: DefinedMemoryIndex) -> Memory {
@@ -908,7 +906,7 @@ impl SharedMemory {
     /// [`ResourceLimiter`](crate::ResourceLimiter) is another example of
     /// preventing a memory to grow.
     pub fn grow(&self, delta: u64) -> Result<u64> {
-        match self.vm.grow(delta, None)? {
+        match self.vm.grow(delta)? {
             Some((old_size, _new_size)) => {
                 // For shared memory, the `VMMemoryDefinition` is updated inside
                 // the locked region.

crates/wasmtime/src/runtime/store.rs

@@ -262,6 +262,27 @@ pub enum StoreResourceLimiter<'a> {
 }
 
 impl StoreResourceLimiter<'_> {
+    pub(crate) async fn memory_growing(
+        &mut self,
+        current: usize,
+        desired: usize,
+        maximum: Option<usize>,
+    ) -> Result<bool, Error> {
+        match self {
+            Self::Sync(s) => s.memory_growing(current, desired, maximum),
+            #[cfg(feature = "async")]
+            Self::Async(s) => s.memory_growing(current, desired, maximum).await,
+        }
+    }
+
+    pub(crate) fn memory_grow_failed(&mut self, error: anyhow::Error) -> Result<()> {
+        match self {
+            Self::Sync(s) => s.memory_grow_failed(error),
+            #[cfg(feature = "async")]
+            Self::Async(s) => s.memory_grow_failed(error),
+        }
+    }
+
     pub(crate) async fn table_growing(
         &mut self,
         current: usize,

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

@@ -62,7 +62,7 @@ impl StoreOpaque {
     fn grow_or_collect_gc_heap(&mut self, bytes_needed: Option<u64>) {
         assert!(!self.async_support());
         if let Some(n) = bytes_needed {
-            if unsafe { self.maybe_async_grow_gc_heap(n).is_ok() } {
+            if vm::assert_ready(self.grow_gc_heap(n)).is_ok() {
                 return;
             }
         }
@@ -72,12 +72,7 @@ impl StoreOpaque {
     /// Attempt to grow the GC heap by `bytes_needed` bytes.
     ///
     /// Returns an error if growing the GC heap fails.
-    ///
-    /// # Safety
-    ///
-    /// When async is enabled, it is the caller's responsibility to ensure that
-    /// this is called on a fiber stack.
-    unsafe fn maybe_async_grow_gc_heap(&mut self, bytes_needed: u64) -> Result<()> {
+    async fn grow_gc_heap(&mut self, bytes_needed: u64) -> Result<()> {
         log::trace!("Attempting to grow the GC heap by {bytes_needed} bytes");
         assert!(bytes_needed > 0);
 
@@ -120,8 +115,15 @@ impl StoreOpaque {
         // `VMMemoryDefinition` in the `VMStoreContext` immediately
         // afterwards.
         unsafe {
+            // FIXME(#11409) this is not sound to widen the borrow
+            let (mut limiter, _) = heap
+                .store
+                .traitobj()
+                .as_mut()
+                .resource_limiter_and_store_opaque();
             heap.memory
-                .grow(delta_pages_for_alloc, Some(heap.store.traitobj().as_mut()))?
+                .grow(delta_pages_for_alloc, limiter.as_mut())
+                .await?
                 .ok_or_else(|| anyhow!("failed to grow GC heap"))?;
         }
         heap.store.vm_store_context.gc_heap = heap.memory.vmmemory();
@@ -246,7 +248,7 @@ impl StoreOpaque {
     async fn grow_or_collect_gc_heap_async(&mut self, bytes_needed: Option<u64>) {
         assert!(self.async_support());
         if let Some(bytes_needed) = bytes_needed {
-            if unsafe { self.maybe_async_grow_gc_heap(bytes_needed).is_ok() } {
+            if self.grow_gc_heap(bytes_needed).await.is_ok() {
                 return;
             }
         }

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

@@ -17,7 +17,7 @@ use crate::runtime::vm::{
     GcStore, HostResult, Imports, ModuleRuntimeInfo, SendSyncPtr, VMGlobalKind, VMStore,
     VMStoreRawPtr, VmPtr, VmSafe, WasmFault, catch_unwind_and_record_trap,
 };
-use crate::store::{InstanceId, StoreId, StoreInstanceId, StoreOpaque};
+use crate::store::{InstanceId, StoreId, StoreInstanceId, StoreOpaque, StoreResourceLimiter};
 use alloc::sync::Arc;
 use core::alloc::Layout;
 use core::marker;
@@ -673,15 +673,18 @@ impl Instance {
     /// Returns `None` if memory can't be grown by the specified amount
     /// of pages. Returns `Some` with the old size in bytes if growth was
     /// successful.
-    pub(crate) fn memory_grow(
+    pub(crate) async fn memory_grow(
         mut self: Pin<&mut Self>,
-        store: &mut dyn VMStore,
+        limiter: Option<&mut StoreResourceLimiter<'_>>,
         idx: DefinedMemoryIndex,
         delta: u64,
     ) -> Result<Option<usize>, Error> {
         let memory = &mut self.as_mut().memories_mut()[idx].1;
 
-        let result = unsafe { memory.grow(delta, Some(store)) };
+        // SAFETY: this is the safe wrapper around `Memory::grow` because it
+        // automatically updates the `VMMemoryDefinition` in this instance after
+        // a growth operation below.
+        let result = unsafe { memory.grow(delta, limiter).await };
 
         // Update the state used by a non-shared Wasm memory in case the base
         // pointer and/or the length changed.

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

@@ -220,12 +220,17 @@ fn memory_grow(
     let module = instance.env_module();
     let page_size_log2 = module.memories[module.memory_index(memory_index)].page_size_log2;
 
-    let result = instance
-        .as_mut()
-        .memory_grow(store, memory_index, delta)?
-        .map(|size_in_bytes| AllocationSize(size_in_bytes >> page_size_log2));
+    let (mut limiter, store) = store.resource_limiter_and_store_opaque();
+    let limiter = limiter.as_mut();
+    block_on!(store, async |_store| {
+        let result = instance
+            .as_mut()
+            .memory_grow(limiter, memory_index, delta)
+            .await?
+            .map(|size_in_bytes| AllocationSize(size_in_bytes >> page_size_log2));
 
-    Ok(result)
+        Ok(result)
+    })?
 }
 
 /// A helper structure to represent the return value of a memory or table growth

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

@@ -75,6 +75,7 @@
 //! they should be merged together.
 
 use crate::prelude::*;
+use crate::runtime::store::StoreResourceLimiter;
 use crate::runtime::vm::vmcontext::VMMemoryDefinition;
 #[cfg(has_virtual_memory)]
 use crate::runtime::vm::{HostAlignedByteCount, MmapOffset};
@@ -394,14 +395,14 @@ impl Memory {
     ///
     /// Ensure that the provided Store is not used to get access any Memory
     /// which lives inside it.
-    pub unsafe fn grow(
+    pub async unsafe fn grow(
         &mut self,
         delta_pages: u64,
-        store: Option<&mut dyn VMStore>,
+        limiter: Option<&mut StoreResourceLimiter<'_>>,
     ) -> Result<Option<usize>, Error> {
         let result = match self {
-            Memory::Local(mem) => mem.grow(delta_pages, store)?,
-            Memory::Shared(mem) => mem.grow(delta_pages, store)?,
+            Memory::Local(mem) => mem.grow(delta_pages, limiter).await?,
+            Memory::Shared(mem) => mem.grow(delta_pages)?,
         };
         match result {
             Some((old, _new)) => Ok(Some(old)),
@@ -600,10 +601,10 @@ impl LocalMemory {
     /// the underlying `grow_to` implementation.
     ///
     /// The `store` is used only for error reporting.
-    pub fn grow(
+    pub async fn grow(
         &mut self,
         delta_pages: u64,
-        mut store: Option<&mut dyn VMStore>,
+        mut limiter: Option<&mut StoreResourceLimiter<'_>>,
     ) -> Result<Option<(usize, usize)>, Error> {
         let old_byte_size = self.alloc.byte_size();
 
@@ -634,8 +635,11 @@ impl LocalMemory {
             .and_then(|n| usize::try_from(n).ok());
 
         // Store limiter gets first chance to reject memory_growing.
-        if let Some(store) = &mut store {
-            if !store.memory_growing(old_byte_size, new_byte_size, maximum)? {
+        if let Some(limiter) = &mut limiter {
+            if !limiter
+                .memory_growing(old_byte_size, new_byte_size, maximum)
+                .await?
+            {
                 return Ok(None);
             }
         }
@@ -701,8 +705,8 @@ impl LocalMemory {
                 // report the growth failure to but the error should not be
                 // dropped
                 // (https://github.com/bytecodealliance/wasmtime/issues/4240).
-                if let Some(store) = store {
-                    store.memory_grow_failed(e)?;
+                if let Some(limiter) = limiter {
+                    limiter.memory_grow_failed(e)?;
                 }
                 Ok(None)
             }

crates/wasmtime/src/runtime/vm/memory/shared_memory.rs

@@ -1,7 +1,7 @@
 use crate::prelude::*;
 use crate::runtime::vm::memory::{LocalMemory, MmapMemory, validate_atomic_addr};
 use crate::runtime::vm::parking_spot::{ParkingSpot, Waiter};
-use crate::runtime::vm::{Memory, VMMemoryDefinition, VMStore, WaitResult};
+use crate::runtime::vm::{self, Memory, VMMemoryDefinition, WaitResult};
 use std::cell::RefCell;
 use std::ops::Range;
 use std::ptr::NonNull;
@@ -68,13 +68,11 @@ impl SharedMemory {
     }
 
     /// Same as `RuntimeLinearMemory::grow`, except with `&self`.
-    pub fn grow(
-        &self,
-        delta_pages: u64,
-        store: Option<&mut dyn VMStore>,
-    ) -> Result<Option<(usize, usize)>, Error> {
+    pub fn grow(&self, delta_pages: u64) -> Result<Option<(usize, usize)>, Error> {
         let mut memory = self.0.memory.write().unwrap();
-        let result = memory.grow(delta_pages, store)?;
+        // Without a limiter being passed in this shouldn't have an await point,
+        // so it should be safe to assert that it's ready.
+        let result = vm::assert_ready(memory.grow(delta_pages, None))?;
         if let Some((_old_size_in_bytes, new_size_in_bytes)) = result {
             // Store the new size to the `VMMemoryDefinition` for JIT-generated
             // code (and runtime functions) to access. No other code can be

crates/wasmtime/src/runtime/vm/memory/shared_memory_disabled.rs

@@ -1,6 +1,6 @@
 use crate::prelude::*;
 use crate::runtime::vm::memory::LocalMemory;
-use crate::runtime::vm::{VMMemoryDefinition, VMStore, WaitResult};
+use crate::runtime::vm::{VMMemoryDefinition, WaitResult};
 use core::ops::Range;
 use core::ptr::NonNull;
 use core::time::Duration;
@@ -26,11 +26,7 @@ impl SharedMemory {
         match *self {}
     }
 
-    pub fn grow(
-        &self,
-        _delta_pages: u64,
-        _store: Option<&mut dyn VMStore>,
-    ) -> Result<Option<(usize, usize)>> {
+    pub fn grow(&self, _delta_pages: u64) -> Result<Option<(usize, usize)>> {
         match *self {}
     }