Reduce type complexity of InstanceAllocator async functions (#12887)

This is a follow-on to #12849 to try to simplify some of the resulting
signatures a bit. Notably the `Result<..., OutOfMemory>` is now packaged
up directly into the output future, so the functions still retain a sort
of "async trait" feel even though they're still incompatible with
`#[async_trait]` (and can't be defined with that anyway).

Author: alexcrichton

crates/wasmtime/src/runtime.rs

@@ -26,6 +26,11 @@
 // situation should be pretty rare though.
 #![warn(clippy::cast_possible_truncation)]
 
+use crate::prelude::*;
+use core::marker;
+use core::pin::Pin;
+use core::task::{Context, Poll};
+
 #[cfg(feature = "component-model-async")]
 mod bug;
 
@@ -127,6 +132,43 @@ pub use coredump::*;
 #[cfg(feature = "wave")]
 mod wave;
 
+/// Helper method to create a future trait object from the future `F` provided.
+///
+/// This requires that the output of `F` is a result where the error can be
+/// created from `OutOfMemory`. This will handle OOM when allocation of `F` on
+/// the heap fails.
+fn box_future<'a, F, T, E>(future: F) -> Pin<Box<dyn Future<Output = Result<T, E>> + Send + 'a>>
+where
+    F: Future<Output = Result<T, E>> + Send + 'a,
+    T: 'a,
+    E: From<OutOfMemory> + 'a,
+{
+    if let Ok(future) = try_new::<Box<F>>(future) {
+        return Pin::from(future);
+    }
+
+    // Use a custom guaranteed-zero-size struct to implement a future that
+    // returns an OOM error which satisfies the type signature of this function.
+    struct OomFuture<F, T, E>(marker::PhantomData<fn() -> (T, F, E)>);
+
+    impl<F, T, E> Future for OomFuture<F, T, E>
+    where
+        E: From<OutOfMemory>,
+    {
+        type Output = Result<T, E>;
+
+        fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
+            Poll::Ready(Err(OutOfMemory::new(size_of::<F>()).into()))
+        }
+    }
+
+    // Zero-size allocations don't actually allocate memory with a `Box`, so
+    // it's ok to use the standard `Box::pin` here and not worry about OOM.
+    let future = OomFuture::<F, T, E>(marker::PhantomData);
+    assert_eq!(size_of_val(&future), 0);
+    Box::pin(future)
+}
+
 fn _assertions_runtime() {
     use crate::_assert_send_and_sync;
 

crates/wasmtime/src/runtime/store.rs

@@ -1958,7 +1958,7 @@ impl StoreOpaque {
 
             let (mem_alloc_index, mem) = engine
                 .allocator()
-                .allocate_memory(&mut request, &mem_ty, None)?
+                .allocate_memory(&mut request, &mem_ty, None)
                 .await?;
 
             // Then, allocate the actual GC heap, passing in that memory

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

@@ -176,10 +176,7 @@ unsafe impl InstanceAllocator for SingleMemoryInstance<'_> {
         request: &'a mut InstanceAllocationRequest<'b, 'c>,
         ty: &'a wasmtime_environ::Memory,
         memory_index: Option<DefinedMemoryIndex>,
-    ) -> Result<
-        Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>>,
-        OutOfMemory,
-    > {
+    ) -> Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>> {
         if cfg!(debug_assertions) {
             let module = request.runtime_info.env_module();
             let offsets = request.runtime_info.offsets();
@@ -188,12 +185,12 @@ unsafe impl InstanceAllocator for SingleMemoryInstance<'_> {
         }
 
         match self.preallocation {
-            Some(shared_memory) => Ok(Box::into_pin(try_new::<Box<_>>(async move {
+            Some(shared_memory) => crate::runtime::box_future(async move {
                 Ok((
                     MemoryAllocationIndex::default(),
                     shared_memory.clone().as_memory(),
                 ))
-            })?)),
+            }),
             None => self.ondemand.allocate_memory(request, ty, memory_index),
         }
     }
@@ -215,10 +212,7 @@ unsafe impl InstanceAllocator for SingleMemoryInstance<'_> {
         req: &'a mut InstanceAllocationRequest<'b, 'c>,
         ty: &'a wasmtime_environ::Table,
         table_index: DefinedTableIndex,
-    ) -> Result<
-        Pin<Box<dyn Future<Output = Result<(TableAllocationIndex, Table)>> + Send + 'a>>,
-        OutOfMemory,
-    > {
+    ) -> Pin<Box<dyn Future<Output = Result<(TableAllocationIndex, Table)>> + Send + 'a>> {
         self.ondemand.allocate_table(req, ty, table_index)
     }
 

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

@@ -193,10 +193,7 @@ pub unsafe trait InstanceAllocator: Send + Sync {
         request: &'a mut InstanceAllocationRequest<'b, 'c>,
         ty: &'a wasmtime_environ::Memory,
         memory_index: Option<DefinedMemoryIndex>,
-    ) -> Result<
-        Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>>,
-        OutOfMemory,
-    >;
+    ) -> Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>>;
 
     /// Deallocate an instance's previously allocated memory.
     ///
@@ -221,10 +218,7 @@ pub unsafe trait InstanceAllocator: Send + Sync {
         req: &'a mut InstanceAllocationRequest<'b, 'c>,
         table: &'a wasmtime_environ::Table,
         table_index: DefinedTableIndex,
-    ) -> Result<
-        Pin<Box<dyn Future<Output = Result<(TableAllocationIndex, Table)>> + Send + 'a>>,
-        OutOfMemory,
-    >;
+    ) -> Pin<Box<dyn Future<Output = Result<(TableAllocationIndex, Table)>> + Send + 'a>>;
 
     /// Deallocate an instance's previously allocated table.
     ///
@@ -404,9 +398,9 @@ impl dyn InstanceAllocator + '_ {
 
     /// Allocate the memories for the given instance allocation request, pushing
     /// them into `memories`.
-    async fn allocate_memories<'a, 'b: 'a, 'c: 'a>(
-        &'a self,
-        request: &'a mut InstanceAllocationRequest<'b, 'c>,
+    async fn allocate_memories(
+        &self,
+        request: &mut InstanceAllocationRequest<'_, '_>,
         memories: &mut TryPrimaryMap<DefinedMemoryIndex, (MemoryAllocationIndex, Memory)>,
     ) -> Result<()> {
         let module = request.runtime_info.env_module();
@@ -422,7 +416,7 @@ impl dyn InstanceAllocator + '_ {
                 .expect("should be a defined memory since we skipped imported ones");
 
             let memory = self
-                .allocate_memory(request, ty, Some(memory_index))?
+                .allocate_memory(request, ty, Some(memory_index))
                 .await?;
             memories.push(memory)?;
         }
@@ -457,9 +451,9 @@ impl dyn InstanceAllocator + '_ {
 
     /// Allocate tables for the given instance allocation request, pushing them
     /// into `tables`.
-    async fn allocate_tables<'a, 'b: 'a, 'c: 'a>(
-        &'a self,
-        request: &'a mut InstanceAllocationRequest<'b, 'c>,
+    async fn allocate_tables(
+        &self,
+        request: &mut InstanceAllocationRequest<'_, '_>,
         tables: &mut TryPrimaryMap<DefinedTableIndex, (TableAllocationIndex, Table)>,
     ) -> Result<()> {
         let module = request.runtime_info.env_module();
@@ -474,7 +468,7 @@ impl dyn InstanceAllocator + '_ {
                 .defined_table_index(index)
                 .expect("should be a defined table since we skipped imported ones");
 
-            let table = self.allocate_table(request, table, def_index)?.await?;
+            let table = self.allocate_table(request, table, def_index).await?;
             tables.push(table)?;
         }
 

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

@@ -115,16 +115,13 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
         request: &'a mut InstanceAllocationRequest<'b, 'c>,
         ty: &'a wasmtime_environ::Memory,
         memory_index: Option<DefinedMemoryIndex>,
-    ) -> Result<
-        Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>>,
-        OutOfMemory,
-    > {
+    ) -> Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>> {
         let creator = self
             .mem_creator
             .as_deref()
             .unwrap_or_else(|| &DefaultMemoryCreator);
 
-        Ok(Box::into_pin(try_new::<Box<_>>(async move {
+        crate::runtime::box_future(async move {
             let image = if let Some(memory_index) = memory_index {
                 request.runtime_info.memory_image(memory_index)?
             } else {
@@ -141,7 +138,7 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
             )
             .await?;
             Ok((allocation_index, memory))
-        })?))
+        })
     }
 
     unsafe fn deallocate_memory(
@@ -159,11 +156,8 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
         request: &'a mut InstanceAllocationRequest<'b, 'c>,
         ty: &'a wasmtime_environ::Table,
         _table_index: DefinedTableIndex,
-    ) -> Result<
-        Pin<Box<dyn Future<Output = Result<(TableAllocationIndex, Table)>> + Send + 'a>>,
-        OutOfMemory,
-    > {
-        Ok(Box::into_pin(try_new::<Box<_>>(async move {
+    ) -> Pin<Box<dyn Future<Output = Result<(TableAllocationIndex, Table)>> + Send + 'a>> {
+        crate::runtime::box_future(async move {
             let allocation_index = TableAllocationIndex::default();
             let table = Table::new_dynamic(
                 ty,
@@ -172,7 +166,7 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
             )
             .await?;
             Ok((allocation_index, table))
-        })?))
+        })
     }
 
     unsafe fn deallocate_table(

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

@@ -684,11 +684,8 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator {
         request: &'a mut InstanceAllocationRequest<'b, 'c>,
         ty: &'a wasmtime_environ::Memory,
         memory_index: Option<DefinedMemoryIndex>,
-    ) -> Result<
-        Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>>,
-        OutOfMemory,
-    > {
-        Ok(Box::into_pin(try_new::<Box<_>>(async move {
+    ) -> Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>> {
+        crate::runtime::box_future(async move {
             async {
                 // FIXME(rust-lang/rust#145127) this should ideally use a version of
                 // `with_flush_and_retry` but adapted for async closures instead of only
@@ -712,7 +709,7 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator {
             .inspect(|_| {
                 self.live_memories.fetch_add(1, Ordering::Relaxed);
             })
-        })?))
+        })
     }
 
     unsafe fn deallocate_memory(
@@ -758,11 +755,9 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator {
         request: &'a mut InstanceAllocationRequest<'b, 'c>,
         ty: &'a wasmtime_environ::Table,
         _table_index: DefinedTableIndex,
-    ) -> Result<
-        Pin<Box<dyn Future<Output = Result<(super::TableAllocationIndex, Table)>> + Send + 'a>>,
-        OutOfMemory,
-    > {
-        Ok(Box::into_pin(try_new::<Box<_>>(async move {
+    ) -> Pin<Box<dyn Future<Output = Result<(super::TableAllocationIndex, Table)>> + Send + 'a>>
+    {
+        crate::runtime::box_future(async move {
             async {
                 // FIXME: see `allocate_memory` above for comments about duplication
                 // with `with_flush_and_retry`.
@@ -784,7 +779,7 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator {
             .inspect(|_| {
                 self.live_tables.fetch_add(1, Ordering::Relaxed);
             })
-        })?))
+        })
     }
 
     unsafe fn deallocate_table(