Add support for async call hooks

crates/wasmtime/src/func.rs

@@ -272,8 +272,9 @@ macro_rules! generate_wrap_async_func {
         {
             assert!(store.as_context().async_support(), concat!("cannot use `wrap", $num, "_async` without enabling async support on the config"));
             Func::wrap(store, move |mut caller: Caller<'_, T>, $($args: $args),*| {
-                let async_cx = caller.store.as_context_mut().0.async_cx();
+                let async_cx = caller.store.as_context_mut().0.async_cx().expect("Attempt to start async function on dying fiber");
                 let mut future = Pin::from(func(caller, $($args),*));
+
                 match unsafe { async_cx.block_on(future.as_mut()) } {
                     Ok(ret) => ret.into_fallible(),
                     Err(e) => R::fallible_from_trap(e),
@@ -439,7 +440,12 @@ impl Func {
             "cannot use `new_async` without enabling async support in the config"
         );
         Func::new(store, ty, move |mut caller, params, results| {
-            let async_cx = caller.store.as_context_mut().0.async_cx();
+            let async_cx = caller
+                .store
+                .as_context_mut()
+                .0
+                .async_cx()
+                .expect("Attempt to spawn new action on dying fiber");
             let mut future = Pin::from(func(caller, params, results));
             match unsafe { async_cx.block_on(future.as_mut()) } {
                 Ok(Ok(())) => Ok(()),

crates/wasmtime/src/lib.rs

@@ -423,6 +423,8 @@ pub use crate::linker::*;
 pub use crate::memory::*;
 pub use crate::module::{FrameInfo, FrameSymbol, Module};
 pub use crate::r#ref::ExternRef;
+#[cfg(feature = "async")]
+pub use crate::store::CallHookHandler;
 pub use crate::store::{AsContext, AsContextMut, CallHook, Store, StoreContext, StoreContextMut};
 pub use crate::trap::*;
 pub use crate::types::*;

crates/wasmtime/src/linker.rs

@@ -146,7 +146,7 @@ macro_rules! generate_wrap_async_func {
                 ),
             );
             self.func_wrap(module, name, move |mut caller: Caller<'_, T>, $($args: $args),*| {
-                let async_cx = caller.store.as_context_mut().0.async_cx();
+                let async_cx = caller.store.as_context_mut().0.async_cx().expect("Attempt to start async function on dying fiber");
                 let mut future = Pin::from(func(caller, $($args),*));
                 match unsafe { async_cx.block_on(future.as_mut()) } {
                     Ok(ret) => ret.into_fallible(),
@@ -360,7 +360,12 @@ impl<T> Linker<T> {
             "cannot use `func_new_async` without enabling async support in the config"
         );
         self.func_new(module, name, ty, move |mut caller, params, results| {
-            let async_cx = caller.store.as_context_mut().0.async_cx();
+            let async_cx = caller
+                .store
+                .as_context_mut()
+                .0
+                .async_cx()
+                .expect("Attempt to spawn new function on dying fiber");
             let mut future = Pin::from(func(caller, params, results));
             match unsafe { async_cx.block_on(future.as_mut()) } {
                 Ok(Ok(())) => Ok(()),

crates/wasmtime/src/store.rs

@@ -200,7 +200,7 @@ pub struct StoreInner<T> {
     inner: StoreOpaque,
 
     limiter: Option<ResourceLimiterInner<T>>,
-    call_hook: Option<Box<dyn FnMut(&mut T, CallHook) -> Result<(), crate::Trap> + Send + Sync>>,
+    call_hook: Option<CallHookInner<T>>,
     // for comments about `ManuallyDrop`, see `Store::into_data`
     data: ManuallyDrop<T>,
 }
@@ -211,6 +211,21 @@ enum ResourceLimiterInner<T> {
     Async(Box<dyn FnMut(&mut T) -> &mut (dyn crate::ResourceLimiterAsync) + Send + Sync>),
 }
 
+/// An object that can take callbacks when the runtime enters or exits hostcalls.
+#[cfg(feature = "async")]
+#[async_trait::async_trait]
+pub trait CallHookHandler<T>: Send {
+    /// A callback to run when wasmtime is about to enter a host call, or when about to
+    /// exit the hostcall.
+    async fn handle_call_event(&self, t: &mut T, ch: CallHook) -> Result<(), crate::Trap>;
+}
+
+enum CallHookInner<T> {
+    Sync(Box<dyn FnMut(&mut T, CallHook) -> Result<(), crate::Trap> + Send + Sync>),
+    #[cfg(feature = "async")]
+    Async(Box<dyn CallHookHandler<T> + Send + Sync>),
+}
+
 // Forward methods on `StoreOpaque` to also being on `StoreInner<T>`
 impl<T> Deref for StoreInner<T> {
     type Target = StoreOpaque;
@@ -603,6 +618,27 @@ impl<T> Store<T> {
         inner.limiter = Some(ResourceLimiterInner::Async(Box::new(limiter)));
     }
 
+    #[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
+    /// Configures an async function that runs on calls and returns between
+    /// WebAssembly and host code. For the non-async equivalent of this method,
+    /// see [`Store::call_hook`].
+    ///
+    /// The function is passed a [`CallHook`] argument, which indicates which
+    /// state transition the VM is making.
+    ///
+    /// This function's future may return a [`Trap`]. If a trap is returned
+    /// when an import was called, it is immediately raised as-if the host
+    /// import had returned the trap. If a trap is returned after wasm returns
+    /// to the host then the wasm function's result is ignored and this trap is
+    /// returned instead.
+    ///
+    /// After this function returns a trap, it may be called for subsequent
+    /// returns to host or wasm code as the trap propagates to the root call.
+    #[cfg(feature = "async")]
+    pub fn call_hook_async(&mut self, hook: impl CallHookHandler<T> + Send + Sync + 'static) {
+        self.inner.call_hook = Some(CallHookInner::Async(Box::new(hook)));
+    }
+
     /// Configure a function that runs on calls and returns between WebAssembly
     /// and host code.
     ///
@@ -616,12 +652,12 @@ impl<T> Store<T> {
     /// instead.
     ///
     /// After this function returns a trap, it may be called for subsequent returns
-    /// to host or wasm code as the trap propogates to the root call.
+    /// to host or wasm code as the trap propagates to the root call.
     pub fn call_hook(
         &mut self,
         hook: impl FnMut(&mut T, CallHook) -> Result<(), Trap> + Send + Sync + 'static,
     ) {
-        self.inner.call_hook = Some(Box::new(hook));
+        self.inner.call_hook = Some(CallHookInner::Sync(Box::new(hook)));
     }
 
     /// Returns the [`Engine`] that this store is associated with.
@@ -956,10 +992,19 @@ impl<T> StoreInner<T> {
     }
 
     pub fn call_hook(&mut self, s: CallHook) -> Result<(), Trap> {
-        if let Some(hook) = &mut self.call_hook {
-            hook(&mut self.data, s)
-        } else {
-            Ok(())
+        match &mut self.call_hook {
+            Some(CallHookInner::Sync(hook)) => hook(&mut self.data, s),
+
+            #[cfg(feature = "async")]
+            Some(CallHookInner::Async(handler)) => unsafe {
+                Ok(self
+                    .inner
+                    .async_cx()
+                    .ok_or(Trap::new("couldn't grab async_cx for call hook"))?
+                    .block_on(handler.handle_call_event(&mut self.data, s).as_mut())??)
+            },
+
+            None => Ok(()),
         }
     }
 }
@@ -1143,14 +1188,29 @@ impl StoreOpaque {
         panic!("trampoline missing")
     }
 
+    /// Yields the async context, assuming that we are executing on a fiber and
+    /// that fiber is not in the process of dying. This function will return
+    /// None in the latter case (the fiber is dying), and panic if
+    /// `async_support()` is false.
     #[cfg(feature = "async")]
     #[inline]
-    pub fn async_cx(&self) -> AsyncCx {
+    pub fn async_cx(&self) -> Option<AsyncCx> {
         debug_assert!(self.async_support());
-        AsyncCx {
-            current_suspend: self.async_state.current_suspend.get(),
-            current_poll_cx: self.async_state.current_poll_cx.get(),
+
+        let poll_cx_box_ptr = self.async_state.current_poll_cx.get();
+        if poll_cx_box_ptr.is_null() {
+            return None;
+        }
+
+        let poll_cx_inner_ptr = unsafe { *poll_cx_box_ptr };
+        if poll_cx_inner_ptr.is_null() {
+            return None;
         }
+
+        Some(AsyncCx {
+            current_suspend: self.async_state.current_suspend.get(),
+            current_poll_cx: poll_cx_box_ptr,
+        })
     }
 
     pub fn fuel_consumed(&self) -> Option<u64> {
@@ -1214,7 +1274,11 @@ impl StoreOpaque {
         // 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().block_on(Pin::new_unchecked(&mut future)) }
+        unsafe {
+            self.async_cx()
+                .expect("attempted to pull async context during shutdown")
+                .block_on(Pin::new_unchecked(&mut future))
+        }
     }
 
     fn add_fuel(&mut self, fuel: u64) -> Result<()> {
@@ -1649,22 +1713,15 @@ unsafe impl<T> wasmtime_runtime::Store 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(feature = "async")]
-        let async_cx = if self.async_support() {
-            Some(self.async_cx())
-        } else {
-            None
-        };
         match self.limiter {
             Some(ResourceLimiterInner::Sync(ref mut limiter)) => {
                 Ok(limiter(&mut self.data).memory_growing(current, desired, maximum))
             }
             #[cfg(feature = "async")]
             Some(ResourceLimiterInner::Async(ref mut limiter)) => unsafe {
-                Ok(async_cx
+                Ok(self
+                    .inner
+                    .async_cx()
                     .expect("ResourceLimiterAsync requires async Store")
                     .block_on(
                         limiter(&mut self.data)
@@ -1700,7 +1757,7 @@ unsafe impl<T> wasmtime_runtime::Store for StoreInner<T> {
         // wrap this in an option.
         #[cfg(feature = "async")]
         let async_cx = if self.async_support() {
-            Some(self.async_cx())
+            Some(self.async_cx().unwrap())
         } else {
             None
         };