Fix sync imports calling code that calls async imports in an async context

Author: Arshia001

lib/api/src/backend/sys/entities/function/mod.rs

@@ -505,6 +505,7 @@ impl Function {
     ) -> Result<(), RuntimeError> {
         // Call the trampoline.
         let result = {
+            let store_id = store.objects_mut().id();
             // Safety: the store context is uninstalled before we return, and the
             // store mut is valid for the duration of the call.
             let store_install_guard =
@@ -516,9 +517,13 @@ impl Function {
                 let storeref = store.as_store_ref();
                 let vm_function = self.handle.get(storeref.objects().as_sys());
                 let config = storeref.engine().tunables().vmconfig();
+                let signal_handler = storeref.signal_handler();
                 r = unsafe {
+                    // Safety: This is the intended use-case for StoreContext::pause, as
+                    // documented in the function's doc comments.
+                    let pause_guard = StoreContext::pause(store_id);
                     wasmer_call_trampoline(
-                        store.as_store_ref().signal_handler(),
+                        signal_handler,
                         config,
                         vm_function.anyfunc.as_ptr().as_ref().vmctx,
                         trampoline,

lib/api/src/backend/sys/entities/function/typed.rs

@@ -52,6 +52,7 @@ macro_rules! impl_native_traits {
                     rets_list.as_mut()
                 };
 
+                let store_id = store.objects_mut().id();
                 // Install the store into the store context
                 let store_install_guard = unsafe {
                     StoreContext::ensure_installed(store.as_store_mut().inner as *mut _)
@@ -62,6 +63,9 @@ macro_rules! impl_native_traits {
                     let storeref = store.as_store_ref();
                     let config = storeref.engine().tunables().vmconfig();
                     r = unsafe {
+                        // Safety: This is the intended use-case for StoreContext::pause, as
+                        // documented in the function's doc comments.
+                        let pause_guard = StoreContext::pause(store_id);
                         wasmer_vm::wasmer_call_trampoline(
                             store.as_store_ref().signal_handler(),
                             config,
@@ -190,6 +194,7 @@ macro_rules! impl_native_traits {
                     rets_list.as_mut()
                 };
 
+                let store_id = store.objects_mut().id();
                 // Install the store into the store context
                 let store_install_guard = unsafe {
                     StoreContext::ensure_installed(store.as_store_mut().inner as *mut _)
@@ -200,6 +205,9 @@ macro_rules! impl_native_traits {
                     let storeref = store.as_store_ref();
                     let config = storeref.engine().tunables().vmconfig();
                     r = unsafe {
+                        // Safety: This is the intended use-case for StoreContext::pause, as
+                        // documented in the function's doc comments.
+                        let pause_guard = StoreContext::pause(store_id);
                         wasmer_vm::wasmer_call_trampoline(
                             store.as_store_ref().signal_handler(),
                             config,

lib/api/src/entities/store/context.rs

@@ -96,6 +96,12 @@ pub(crate) struct StoreAsyncGuardWrapper {
     pub(crate) guard: *mut LocalRwLockWriteGuard<Box<StoreInner>>,
 }
 
+pub(crate) struct StorePtrPauseGuard {
+    store_id: StoreId,
+    ptr: *mut StoreInner,
+    ref_count_decremented: bool,
+}
+
 #[cfg(feature = "experimental-async")]
 pub(crate) enum GetStoreAsyncGuardResult {
     Ok(StoreAsyncGuardWrapper),
@@ -170,13 +176,56 @@ impl StoreContext {
     pub(crate) unsafe fn ensure_installed(store_ptr: *mut StoreInner) -> StoreInstallGuard {
         let store_id = unsafe { store_ptr.as_ref().unwrap().objects.id() };
         if Self::is_active(store_id) {
+            let current_ptr = STORE_CONTEXT_STACK.with(|cell| {
+                let stack = cell.borrow();
+                unsafe { stack.last().unwrap().entry.get().as_ref().unwrap().as_ptr() }
+            });
+            assert_eq!(store_ptr, current_ptr, "Store context pointer mismatch");
             StoreInstallGuard::NotInstalled
         } else {
             Self::install(store_id, StoreContextEntry::Sync(store_ptr));
             StoreInstallGuard::Installed(store_id)
         }
     }
 
+    /// "Pause" one borrow of the store context.
+    ///
+    /// # Safety
+    /// Code must ensure it does not use the StorePtrWrapper or
+    /// StoreAsyncGuardWrapper that it owns, or any StoreRef/StoreMut
+    /// derived from them, while the store context is paused.
+    ///
+    /// The safe, correct use-case for this method is to
+    /// pause the store context while executing WASM code, which
+    /// cannot use the store context directly. This allows an async
+    /// context to uninstall the store context when suspending if it's
+    /// called from a sync imported function. The imported function
+    /// will have borrowed the store context in its trampoline, which
+    /// will prevent the async context from uninstalling the store.
+    /// However, since the imported function passes a mutable borrow
+    /// of its store into `Function::call`, it will expect the store
+    /// to change before the call returns.
+    pub(crate) unsafe fn pause(id: StoreId) -> StorePtrPauseGuard {
+        STORE_CONTEXT_STACK.with(|cell| {
+            let mut stack = cell.borrow_mut();
+            let top = stack
+                .last_mut()
+                .expect("No store context installed on this thread");
+            assert_eq!(top.id, id, "Mismatched store context access");
+            let ref_count_decremented = if top.borrow_count > 0 {
+                top.borrow_count -= 1;
+                true
+            } else {
+                false
+            };
+            StorePtrPauseGuard {
+                store_id: id,
+                ptr: unsafe { top.entry.get().as_ref().unwrap().as_ptr() },
+                ref_count_decremented,
+            }
+        })
+    }
+
     /// Safety: This method lets you borrow multiple mutable references
     /// to the currently active store context. The caller must ensure that:
     ///   * there is only one mutable reference alive, or
@@ -291,6 +340,9 @@ impl Clone for StorePtrWrapper {
 
 impl Drop for StorePtrWrapper {
     fn drop(&mut self) {
+        if std::thread::panicking() {
+            return;
+        }
         let id = self.as_mut().objects_mut().id();
         STORE_CONTEXT_STACK.with(|cell| {
             let mut stack = cell.borrow_mut();
@@ -306,6 +358,9 @@ impl Drop for StorePtrWrapper {
 #[cfg(feature = "experimental-async")]
 impl Drop for StoreAsyncGuardWrapper {
     fn drop(&mut self) {
+        if std::thread::panicking() {
+            return;
+        }
         let id = unsafe { self.guard.as_ref().unwrap().objects.id() };
         STORE_CONTEXT_STACK.with(|cell| {
             let mut stack = cell.borrow_mut();
@@ -323,12 +378,28 @@ impl Drop for StoreInstallGuard {
         if let Self::Installed(store_id) = self {
             STORE_CONTEXT_STACK.with(|cell| {
                 let mut stack = cell.borrow_mut();
-                let top = stack.pop().expect("Store context stack underflow");
-                assert_eq!(top.id, *store_id, "Mismatched store context uninstall");
-                assert_eq!(
-                    top.borrow_count, 0,
-                    "Cannot uninstall store context while it is still borrowed"
-                );
+                match (stack.pop(), std::thread::panicking()) {
+                    (Some(top), false) => {
+                        assert_eq!(top.id, *store_id, "Mismatched store context uninstall");
+                        assert_eq!(
+                            top.borrow_count, 0,
+                            "Cannot uninstall store context while it is still borrowed"
+                        );
+                    }
+                    (Some(top), true) => {
+                        // If we're panicking and there's a store ID mismatch, just
+                        // put the store back in the hope that its own install guard
+                        // take care of uninstalling it later.
+                        if top.id != *store_id {
+                            stack.push(top);
+                        }
+                    }
+                    (None, false) => panic!("Store context stack underflow"),
+                    (None, true) => {
+                        // Nothing to do if we're panicking; panics can put the context
+                        // in an invalid state, and we don't to cause another panic here.
+                    }
+                }
             })
         }
     }
@@ -338,12 +409,51 @@ impl Drop for ForcedStoreInstallGuard {
     fn drop(&mut self) {
         STORE_CONTEXT_STACK.with(|cell| {
             let mut stack = cell.borrow_mut();
-            let top = stack.pop().expect("Store context stack underflow");
-            assert_eq!(top.id, self.store_id, "Mismatched store context uninstall");
+            match (stack.pop(), std::thread::panicking()) {
+                (Some(top), false) => {
+                    assert_eq!(top.id, self.store_id, "Mismatched store context uninstall");
+                    assert_eq!(
+                        top.borrow_count, 0,
+                        "Cannot uninstall store context while it is still borrowed"
+                    );
+                }
+                (Some(top), true) => {
+                    // If we're panicking and there's a store ID mismatch, just
+                    // put the store back in the hope that its own install guard
+                    // take care of uninstalling it later.
+                    if top.id != self.store_id {
+                        stack.push(top);
+                    }
+                }
+                (None, false) => panic!("Store context stack underflow"),
+                (None, true) => {
+                    // Nothing to do if we're panicking; panics can put the context
+                    // in an invalid state, and we don't to cause another panic here.
+                }
+            }
+        })
+    }
+}
+
+impl Drop for StorePtrPauseGuard {
+    fn drop(&mut self) {
+        if std::thread::panicking() {
+            return;
+        }
+        STORE_CONTEXT_STACK.with(|cell| {
+            let mut stack = cell.borrow_mut();
+            let top = stack
+                .last_mut()
+                .expect("No store context installed on this thread");
+            assert_eq!(top.id, self.store_id, "Mismatched store context access");
             assert_eq!(
-                top.borrow_count, 0,
-                "Cannot uninstall store context while it is still borrowed"
+                unsafe { top.entry.get().as_ref().unwrap() }.as_ptr(),
+                self.ptr,
+                "Mismatched store context access"
             );
+            if self.ref_count_decremented {
+                top.borrow_count += 1;
+            }
         })
     }
 }

lib/api/src/error.rs

@@ -269,7 +269,7 @@ impl RuntimeError {
     pub(crate) fn from_dyn(err: Box<dyn std::error::Error + Send + Sync>) -> Self {
         match err.downcast::<Self>() {
             Ok(runtime_error) => *runtime_error,
-            Err(error) => Trap::user(error).into(),
+            Err(error) => Trap::user(error),
         }
     }
 }

lib/api/tests/jspi_async.rs

@@ -1,12 +1,12 @@
 #![cfg(feature = "experimental-async")]
 
-use std::sync::OnceLock;
+use std::{cell::RefCell, sync::OnceLock};
 
 use anyhow::Result;
 use futures::future;
 use wasmer::{
-    AsyncFunctionEnvMut, Function, FunctionEnv, FunctionType, Instance, Module, Store, StoreAsync,
-    Type, TypedFunction, Value, imports,
+    AsyncFunctionEnvMut, Function, FunctionEnv, FunctionEnvMut, FunctionType, Instance, Module,
+    Store, StoreAsync, Type, TypedFunction, Value, imports,
 };
 use wasmer_vm::TrapCode;
 
@@ -227,3 +227,79 @@ fn cannot_yield_when_not_in_async_context() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn nested_async_in_sync() -> Result<()> {
+    const WAT: &str = r#"
+    (module
+        (import "env" "sync" (func $sync (result i32)))
+        (import "env" "async" (func $async (result i32)))
+        (func (export "entry") (result i32)
+            call $sync
+        )
+        (func (export "inner_async") (result i32)
+            call $async
+        )
+    )
+    "#;
+    let wasm = wat::parse_str(WAT).expect("valid WAT module");
+
+    let mut store = Store::default();
+    let module = Module::new(&store, wasm)?;
+
+    struct Env {
+        inner_async: RefCell<Option<wasmer::TypedFunction<(), i32>>>,
+    }
+    let env = FunctionEnv::new(
+        &mut store,
+        Env {
+            inner_async: RefCell::new(None),
+        },
+    );
+
+    let sync = Function::new_typed_with_env(&mut store, &env, |mut env: FunctionEnvMut<Env>| {
+        let (env, mut store) = env.data_and_store_mut();
+        env.inner_async
+            .borrow()
+            .as_ref()
+            .expect("inner_async function to be set")
+            .call(&mut store)
+            .expect("inner async call to succeed")
+    });
+
+    let async_ = Function::new_typed_async(&mut store, async || {
+        tokio::task::yield_now().await;
+        42
+    });
+
+    let imports = imports! {
+        "env" => {
+            "sync" => sync,
+            "async" => async_,
+        }
+    };
+
+    let instance = Instance::new(&mut store, &module, &imports)?;
+
+    let inner_async = instance
+        .exports
+        .get_typed_function::<(), i32>(&mut store, "inner_async")
+        .unwrap();
+    env.as_mut(&mut store)
+        .inner_async
+        .borrow_mut()
+        .replace(inner_async);
+
+    let entry = instance
+        .exports
+        .get_typed_function::<(), i32>(&mut store, "entry")?;
+    let result = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap()
+        .block_on(entry.call_async(&store.into_async()))?;
+
+    assert_eq!(result, 42);
+
+    Ok(())
+}

tests/compilers/typed_functions.rs

@@ -22,7 +22,7 @@ fn long_f(a: u32, b: u32, c: u32, d: u32, e: u32, f: u16, g: u64, h: u64, i: u16
         + a as u64 * 1000000000
 }
 
-fn long_f_dynamic(values: &[Value]) -> Result<Vec<Value>, RuntimeError> {
+fn long_f_dynamic(values: &[Value]) -> DynamicFunctionResult {
     Ok(vec![Value::I64(
         values[9].unwrap_i32() as i64
             + values[8].unwrap_i32() as i64 * 10