Refactor how block_on is called and used

Author: alexcrichton

crates/wasmtime/src/runtime/store.rs

@@ -1126,33 +1126,9 @@ impl<T> StoreInner<T> {
             CallHookInner::Sync(hook) => hook((&mut *self).as_context_mut(), s),
 
             #[cfg(all(feature = "async", feature = "call-hook"))]
-            CallHookInner::Async(handler) => unsafe {
-                #[cfg(feature = "component-model-async")]
-                {
-                    let async_cx = crate::component::concurrent::AsyncCx::try_new(&mut self.inner)
-                        .ok_or_else(|| anyhow!("couldn't grab async_cx for call hook"))?;
-
-                    async_cx
-                        .block_on(
-                            handler
-                                .handle_call_event((&mut *self).as_context_mut(), s)
-                                .as_mut(),
-                            None,
-                        )?
-                        .0
-                }
-                #[cfg(not(feature = "component-model-async"))]
-                {
-                    self.inner
-                        .async_cx()
-                        .ok_or_else(|| anyhow!("couldn't grab async_cx for call hook"))?
-                        .block_on(
-                            handler
-                                .handle_call_event((&mut *self).as_context_mut(), s)
-                                .as_mut(),
-                        )?
-                }
-            },
+            CallHookInner::Async(handler) => {
+                self.block_on(|store| handler.handle_call_event(StoreContextMut(store), s))?
+            }
 
             CallHookInner::ForceTypeParameterToBeUsed { uninhabited, .. } => {
                 let _ = s;
@@ -2047,36 +2023,12 @@ unsafe impl<T> crate::runtime::vm::VMStore for StoreInner<T> {
                 limiter(&mut self.data).memory_growing(current, desired, maximum)
             }
             #[cfg(feature = "async")]
-            Some(ResourceLimiterInner::Async(ref mut limiter)) => unsafe {
-                #[cfg(feature = "component-model-async")]
-                {
-                    _ = limiter;
-                    let async_cx =
-                        crate::component::concurrent::AsyncCx::new(&mut (&mut *self).inner);
-                    let Some(ResourceLimiterInner::Async(ref mut limiter)) = self.limiter else {
-                        unreachable!();
-                    };
-                    async_cx
-                        .block_on(
-                            limiter(&mut self.data)
-                                .memory_growing(current, desired, maximum)
-                                .as_mut(),
-                            None,
-                        )?
-                        .0
-                }
-                #[cfg(not(feature = "component-model-async"))]
-                {
-                    self.inner
-                        .async_cx()
-                        .expect("ResourceLimiterAsync requires async Store")
-                        .block_on(
-                            limiter(&mut self.data)
-                                .memory_growing(current, desired, maximum)
-                                .as_mut(),
-                        )?
-                }
-            },
+            Some(ResourceLimiterInner::Async(_)) => self.block_on(|store| {
+                let Some(ResourceLimiterInner::Async(limiter)) = &mut store.limiter else {
+                    unreachable!();
+                };
+                limiter(&mut store.data).memory_growing(current, desired, maximum)
+            })?,
             None => Ok(true),
         }
     }
@@ -2103,50 +2055,17 @@ unsafe impl<T> crate::runtime::vm::VMStore for StoreInner<T> {
         desired: usize,
         maximum: Option<usize>,
     ) -> Result<bool, anyhow::Error> {
-        // Need to borrow async_cx before the mut borrow of the limiter.
-        // self.async_cx() panicks when used with a non-async store, so
-        // wrap this in an option.
-        #[cfg(all(feature = "async", not(feature = "component-model-async")))]
-        let async_cx = if self.async_support()
-            && matches!(self.limiter, Some(ResourceLimiterInner::Async(_)))
-        {
-            Some(self.async_cx().unwrap())
-        } else {
-            None
-        };
-
         match self.limiter {
             Some(ResourceLimiterInner::Sync(ref mut limiter)) => {
                 limiter(&mut self.data).table_growing(current, desired, maximum)
             }
             #[cfg(feature = "async")]
-            Some(ResourceLimiterInner::Async(ref mut limiter)) => unsafe {
-                #[cfg(feature = "component-model-async")]
-                {
-                    _ = limiter;
-                    let async_cx =
-                        crate::component::concurrent::AsyncCx::new(&mut (&mut *self).inner);
-                    let Some(ResourceLimiterInner::Async(ref mut limiter)) = self.limiter else {
-                        unreachable!();
-                    };
-                    async_cx
-                        .block_on(
-                            limiter(&mut self.data)
-                                .table_growing(current, desired, maximum)
-                                .as_mut(),
-                            None,
-                        )?
-                        .0
-                }
-                #[cfg(not(feature = "component-model-async"))]
-                {
-                    async_cx
-                        .expect("ResourceLimiterAsync requires async Store")
-                        .block_on(
-                            limiter(&mut self.data).table_growing(current, desired, maximum),
-                        )?
-                }
-            },
+            Some(ResourceLimiterInner::Async(_)) => self.block_on(|store| {
+                let Some(ResourceLimiterInner::Async(limiter)) = &mut store.limiter else {
+                    unreachable!();
+                };
+                limiter(&mut store.data).table_growing(current, desired, maximum)
+            })?,
             None => Ok(true),
         }
     }
@@ -2213,11 +2132,7 @@ unsafe impl<T> crate::runtime::vm::VMStore for StoreInner<T> {
                         // to clean up this fiber. Do so by raising a trap which will
                         // abort all wasm and get caught on the other side to clean
                         // things up.
-                        unsafe {
-                            self.async_cx()
-                                .expect("attempted to pull async context during shutdown")
-                                .block_on(future)?
-                        }
+                        self.block_on(|_| future)?;
                         delta
                     }
                 };

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

@@ -203,29 +203,14 @@ impl<T> StoreInner<T> {
     /// This only works on async futures and stores, and assumes that we're
     /// executing on a fiber. This will yield execution back to the caller once.
     pub fn async_yield_impl(&mut self) -> Result<()> {
-        use crate::runtime::vm::Yield;
-
-        let mut future = Yield::new();
-
         // When control returns, we have a `Result<()>` passed
         // in from the host fiber. If this finished successfully then
         // we were resumed normally via a `poll`, so keep going.  If
         // the future was dropped while we were yielded, then we need
         // to clean up this fiber. Do so by raising a trap which will
         // abort all wasm and get caught on the other side to clean
         // things up.
-        #[cfg(feature = "component-model-async")]
-        unsafe {
-            let async_cx = crate::component::concurrent::AsyncCx::new(&mut (&mut *self).inner);
-            async_cx.block_on(Pin::new_unchecked(&mut future), None)?.0;
-            Ok(())
-        }
-        #[cfg(not(feature = "component-model-async"))]
-        unsafe {
-            self.async_cx()
-                .expect("attempted to pull async context during shutdown")
-                .block_on(Pin::new_unchecked(&mut future))
-        }
+        self.block_on(|_| crate::runtime::vm::Yield::new())
     }
 
     #[cfg(target_has_atomic = "64")]
@@ -237,6 +222,30 @@ impl<T> StoreInner<T> {
         self.epoch_deadline_behavior =
             Some(Box::new(move |_store| Ok(UpdateDeadline::Yield(delta))));
     }
+
+    pub(crate) fn block_on<'a, F: Future + Send>(
+        &'a mut self,
+        mk_future: impl FnOnce(&'a mut Self) -> F,
+    ) -> Result<F::Output> {
+        #[cfg(feature = "component-model-async")]
+        {
+            let async_cx = crate::component::concurrent::AsyncCx::try_new(self)
+                .expect("couldn't create AsyncCx to block on async operation");
+            let future = mk_future(self);
+            let mut future = core::pin::pin!(future);
+            unsafe { Ok(async_cx.block_on(future.as_mut(), None)?.0) }
+        }
+        #[cfg(not(feature = "component-model-async"))]
+        {
+            let cx = self
+                .inner
+                .async_cx()
+                .expect("async_cx is not present to block on async operation");
+            let future = mk_future(self);
+            let mut future = core::pin::pin!(future);
+            unsafe { cx.block_on(future.as_mut()) }
+        }
+    }
 }
 
 #[doc(hidden)]