Make memory growth an async function (bytecodealliance#11460)

* Make memory growth an `async` function

This is an analog of bytecodealliance#11442 but for memories. This had a little more
impact due to memories being hooked into GC operations. Further
refactoring of GC operations to make them safer/more-async is deferred
to a future PR and for now it's "no worse than before". This is another
step towards bytecodealliance#11430 and enables removing a longstanding `unsafe` block
in `runtime/memory.rs` which previously could not be removed.

One semantic change from this is that growth of a shared memory no
longer uses an async limiter. This is done to keep growth of a shared
memory consistent with creation of a shared memory where no limits are
applied. This is due to the cross-store nature of shared memories which
means that we can't tie growth to any one particular store. This
additionally fixes an issue where an rwlock write guard was otherwise
held across a `.await` point which creates a non-`Send` future, closing
a possible soundness hole in Wasmtime.

* Fix threads-disabled build

* Review comments

Author: alexcrichton

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 {}
     }