Repository: bytecodealliance/wasmtime
Pull request #7220 (2023-10-11 18:46:42)

threads: log every `wait` and `notify`
When troubleshooting deadlocks in WebAssembly modules, it is important to understand which `wait` instructions are still pending a `notify`. It would be nice to have some kind of `--warn-deadlock-after=1s` flag available that would poll the parking lot for `wait`s hanging past the time limit, but I realized the real value would be to tie the `wait` instruction (through CLIF) to the original source code, if debug information were available. This did not seem to be entirely feasible, since CLIF loses the original Wasm source context (is this true?) and I was not confident that we would be able to use `addr2line` to map from Wasm instructions to source (e.g., see @cfallin's [issue](https://github.com/gimli-rs/addr2line/issues/265)).

Instead, this change simply logs each valid `wait` and `notify` execution, leaving it to the user to figure out which one is hanging (should not be too difficult) and how to map this back to their source code (more difficult).

<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->



Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 686bcac67..85a492253 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3523,6 +3523,7 @@ dependencies = [
 name = "wasmtime-fiber"
 version = "15.0.0"
 dependencies = [
+ "anyhow",
  "backtrace",
  "cc",
  "cfg-if",
diff --git a/cranelift/codegen/src/isa/riscv64/inst/encode.rs b/cranelift/codegen/src/isa/riscv64/inst/encode.rs
index fae2f929c..c169c11ec 100644
--- a/cranelift/codegen/src/isa/riscv64/inst/encode.rs
+++ b/cranelift/codegen/src/isa/riscv64/inst/encode.rs
@@ -468,7 +468,7 @@ pub fn encode_ciw_type(op: CiwOp, rd: WritableReg, imm: u8) -> u16 {
     let mut imm_field = 0;
     imm_field |= ((imm >> 1) & 1) << 0;
     imm_field |= ((imm >> 0) & 1) << 1;
-    imm_field |= ((imm >> 4) & 7) << 2;
+    imm_field |= ((imm >> 4) & 15) << 2;
     imm_field |= ((imm >> 2) & 3) << 6;
 
     let mut bits = 0;
diff --git a/crates/c-api/include/wasmtime/async.h b/crates/c-api/include/wasmtime/async.h
index dbe4312dc..3919dc3f1 100644
--- a/crates/c-api/include/wasmtime/async.h
+++ b/crates/c-api/include/wasmtime/async.h
@@ -282,6 +282,80 @@ WASM_API_EXTERN wasmtime_call_future_t *wasmtime_instance_pre_instantiate_async(
     wasm_trap_t** trap_ret,
     wasmtime_error_t** error_ret);
 
+/**
+ * A callback to get the top of the stack address and the length of the stack,
+ * excluding guard pages.
+ *
+ * For more information about the parameters see the Rust documentation at
+ * https://docs.wasmtime.dev/api/wasmtime/trait.StackMemory.html
+ */
+typedef uint8_t *(*wasmtime_stack_memory_get_callback_t)(
+    void *env,
+    size_t *out_len);
+
+/**
+ * A Stack instance created from a #wasmtime_new_stack_memory_callback_t.
+ *
+ * For more information see the Rust documentation at
+ * https://docs.wasmtime.dev/api/wasmtime/trait.StackMemory.html
+ */
+typedef struct {
+  /// User provided value to be passed to get_memory and grow_memory
+  void *env;
+  /// Callback to get the memory and size of this LinearMemory
+  wasmtime_stack_memory_get_callback_t get_stack_memory;
+  /// An optional finalizer for env
+  void (*finalizer)(void*);
+} wasmtime_stack_memory_t;
+
+/**
+ * A callback to create a new StackMemory from the specified parameters.
+ *
+ * The result should be written to `stack_ret` and wasmtime will own the values written 
+ * into that struct.
+ *
+ * This callback must be thread-safe.
+ *
+ * For more information about the parameters see the Rust documentation at
+ * https://docs.wasmtime.dev/api/wasmtime/trait.StackCreator.html#tymethod.new_stack
+ */
+typedef wasmtime_error_t *(*wasmtime_new_stack_memory_callback_t)(
+    void *env,
+    size_t size,
+    wasmtime_stack_memory_t *stack_ret);
+
+/**
+ * A representation of custom stack creator.
+ *
+ * For more information see the Rust documentation at
+ * https://docs.wasmtime.dev/api/wasmtime/trait.StackCreator.html
+ */
+typedef struct {
+  /// User provided value to be passed to new_stack
+  void* env;
+  /// The callback to create a new stack, must be thread safe
+  wasmtime_new_stack_memory_callback_t new_stack;
+  /// An optional finalizer for env.
+  void (*finalizer)(void*);
+} wasmtime_stack_creator_t;
+
+/**
+ * Sets a custom stack creator.
+ *
+ * Custom memory creators are used when creating creating async instance stacks for
+ * the on-demand instance allocation strategy.
+ *
+ * The config does **not** take ownership of the #wasmtime_stack_creator_t passed in, but 
+ * instead copies all the values in the struct.
+ *
+ * For more information see the Rust documentation at
+ * https://docs.wasmtime.dev/api/wasmtime/struct.Config.html#method.with_host_stack
+ */
+WASM_API_EXTERN void wasmtime_config_host_stack_creator_set(
+    wasm_config_t*,
+    wasmtime_memory_creator_t*);
+
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif
diff --git a/crates/c-api/src/async.rs b/crates/c-api/src/async.rs
index 23572db32..60cf7d308 100644
--- a/crates/c-api/src/async.rs
+++ b/crates/c-api/src/async.rs
@@ -1,11 +1,15 @@
 use std::ffi::c_void;
 use std::future::Future;
 use std::mem::{self, MaybeUninit};
+use std::ops::Range;
 use std::pin::Pin;
+use std::sync::Arc;
 use std::task::{Context, Poll};
 use std::{ptr, str};
 
-use wasmtime::{AsContextMut, Caller, Func, Instance, Result, Trap, Val};
+use wasmtime::{
+    AsContextMut, Caller, Func, Instance, Result, StackCreator, StackMemory, Trap, Val,
+};
 
 use crate::{
     bad_utf8, handle_result, to_str, translate_args, wasm_config_t, wasm_functype_t, wasm_trap_t,
@@ -335,3 +339,92 @@ pub extern "C" fn wasmtime_instance_pre_instantiate_async<'a>(
     ));
     Box::new(crate::wasmtime_call_future_t { underlying: fut })
 }
+
+pub type wasmtime_stack_memory_get_callback_t =
+    extern "C" fn(env: *mut std::ffi::c_void, out_len: &mut usize) -> *mut u8;
+
+#[repr(C)]
+pub struct wasmtime_stack_memory_t {
+    env: *mut std::ffi::c_void,
+    get_stack_memory: wasmtime_stack_memory_get_callback_t,
+    finalizer: Option<extern "C" fn(arg1: *mut std::ffi::c_void)>,
+}
+
+struct CHostStackMemory {
+    foreign: crate::ForeignData,
+    get_memory: wasmtime_stack_memory_get_callback_t,
+}
+unsafe impl Send for CHostStackMemory {}
+unsafe impl Sync for CHostStackMemory {}
+unsafe impl StackMemory for CHostStackMemory {
+    fn top(&self) -> *mut u8 {
+        let mut len = 0;
+        let cb = self.get_memory;
+        cb(self.foreign.data, &mut len)
+    }
+    fn range(&self) -> Range<usize> {
+        let mut len = 0;
+        let cb = self.get_memory;
+        let top = cb(self.foreign.data, &mut len);
+        let base = unsafe { top.sub(len) as usize };
+        base..base + len
+    }
+}
+
+pub type wasmtime_new_stack_memory_callback_t = extern "C" fn(
+    env: *mut std::ffi::c_void,
+    size: usize,
+    stack_ret: &mut wasmtime_stack_memory_t,
+) -> Option<Box<wasmtime_error_t>>;
+
+#[repr(C)]
+pub struct wasmtime_stack_creator_t {
+    env: *mut std::ffi::c_void,
+    new_stack: wasmtime_new_stack_memory_callback_t,
+    finalizer: Option<extern "C" fn(arg1: *mut std::ffi::c_void)>,
+}
+
+struct CHostStackCreator {
+    foreign: crate::ForeignData,
+    new_stack: wasmtime_new_stack_memory_callback_t,
+}
+unsafe impl Send for CHostStackCreator {}
+unsafe impl Sync for CHostStackCreator {}
+unsafe impl StackCreator for CHostStackCreator {
+    fn new_stack(&self, size: usize) -> Result<Box<dyn wasmtime::StackMemory>> {
+        extern "C" fn panic_callback(_env: *mut std::ffi::c_void, _out_len: &mut usize) -> *mut u8 {
+            panic!("a callback must be set");
+        }
+        let mut out = wasmtime_stack_memory_t {
+            env: ptr::null_mut(),
+            get_stack_memory: panic_callback,
+            finalizer: None,
+        };
+        let cb = self.new_stack;
+        let result = cb(self.foreign.data, size, &mut out);
+        match result {
+            Some(error) => Err((*error).into()),
+            None => Ok(Box::new(CHostStackMemory {
+                foreign: crate::ForeignData {
+                    data: out.env,
+                    finalizer: out.finalizer,
+                },
+                get_memory: out.get_stack_memory,
+            })),
+        }
+    }
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn wasmtime_config_host_stack_creator_set(
+    c: &mut wasm_config_t,
+    creator: &wasmtime_stack_creator_t,
+) {
+    c.config.with_host_stack(Arc::new(CHostStackCreator {
+        foreign: crate::ForeignData {
+            data: creator.env,
+            finalizer: creator.finalizer,
+        },
+        new_stack: creator.new_stack,
+    }));
+}
diff --git a/crates/c-api/src/config.rs b/crates/c-api/src/config.rs
index 76fe4546b..bcf8acde6 100644
--- a/crates/c-api/src/config.rs
+++ b/crates/c-api/src/config.rs
@@ -3,9 +3,9 @@
 #![cfg_attr(not(feature = "cache"), allow(unused_imports))]
 
 use crate::{handle_result, wasm_memorytype_t, wasmtime_error_t};
-use std::mem::MaybeUninit;
 use std::ops::Range;
 use std::os::raw::c_char;
+use std::ptr;
 use std::{ffi::CStr, sync::Arc};
 use wasmtime::{
     Config, LinearMemory, MemoryCreator, OptLevel, ProfilingStrategy, Result, Strategy,
@@ -362,7 +362,25 @@ unsafe impl MemoryCreator for CHostMemoryCreator {
         reserved_size_in_bytes: Option<usize>,
         guard_size_in_bytes: usize,
     ) -> Result<Box<dyn wasmtime::LinearMemory>, String> {
-        let mut memory = MaybeUninit::uninit();
+        extern "C" fn panic_get_callback(
+            _env: *mut std::ffi::c_void,
+            _byte_size: &mut usize,
+            _maximum_byte_size: &mut usize,
+        ) -> *mut u8 {
+            panic!("a callback must be set");
+        }
+        extern "C" fn panic_grow_callback(
+            _env: *mut std::ffi::c_void,
+            _size: usize,
+        ) -> Option<Box<wasmtime_error_t>> {
+            panic!("a callback must be set");
+        }
+        let mut memory = wasmtime_linear_memory_t {
+            env: ptr::null_mut(),
+            get_memory: panic_get_callback,
+            grow_memory: panic_grow_callback,
+            finalizer: None,
+        };
         let cb = self.new_memory;
         let error = cb(
             self.foreign.data,
@@ -371,11 +389,10 @@ unsafe impl MemoryCreator for CHostMemoryCreator {
             maximum.unwrap_or(usize::MAX),
             reserved_size_in_bytes.unwrap_or(0),
             guard_size_in_bytes,
-            memory.as_mut_ptr(),
+            &mut memory,
         );
         match error {
             None => {
-                let memory = unsafe { memory.assume_init() };
                 let foreign = crate::ForeignData {
                     data: memory.env,
                     finalizer: memory.finalizer,
diff --git a/crates/fiber/Cargo.toml b/crates/fiber/Cargo.toml
index 4c221ee1f..555733050 100644
--- a/crates/fiber/Cargo.toml
+++ b/crates/fiber/Cargo.toml
@@ -8,6 +8,7 @@ repository = "https://github.com/bytecodealliance/wasmtime"
 edition.workspace = true
 
 [dependencies]
+anyhow = { workspace = true }
 cfg-if = { workspace = true }
 wasmtime-versioned-export-macros = { workspace = true }
 
diff --git a/crates/fiber/src/lib.rs b/crates/fiber/src/lib.rs
index 79d7bdaa6..5b20157cf 100644
--- a/crates/fiber/src/lib.rs
+++ b/crates/fiber/src/lib.rs
@@ -1,3 +1,4 @@
+use anyhow::Error;
 use std::any::Any;
 use std::cell::Cell;
 use std::io;
@@ -18,7 +19,6 @@ cfg_if::cfg_if! {
 }
 
 /// Represents an execution stack to use for a fiber.
-#[derive(Debug)]
 pub struct FiberStack(imp::FiberStack);
 
 impl FiberStack {
@@ -27,6 +27,11 @@ impl FiberStack {
         Ok(Self(imp::FiberStack::new(size)?))
     }
 
+    /// Creates a new fiber stack of the given size.
+    pub fn from_custom(custom: Box<dyn RuntimeFiberStack>) -> io::Result<Self> {
+        Ok(Self(imp::FiberStack::from_custom(custom)?))
+    }
+
     /// Creates a new fiber stack with the given pointer to the bottom of the
     /// stack plus the byte length of the stack.
     ///
@@ -58,6 +63,23 @@ impl FiberStack {
     }
 }
 
+/// A creator of RuntimeFiberStacks.
+pub unsafe trait RuntimeFiberStackCreator: Send + Sync {
+    /// Creates a new RuntimeFiberStack with the specified size, guard pages should be included.
+    ///
+    /// This is useful to plugin previously allocated memory instead of mmap'ing a new stack for
+    /// every instance.
+    fn new_stack(&self, size: usize) -> Result<Box<dyn RuntimeFiberStack>, Error>;
+}
+
+/// A fiber stack backed by custom memory.
+pub unsafe trait RuntimeFiberStack: Send + Sync {
+    /// The top of the allocated stack.
+    fn top(&self) -> *mut u8;
+    /// The valid range of the stack without guard pages.
+    fn range(&self) -> Range<usize>;
+}
+
 pub struct Fiber<'a, Resume, Yield, Return> {
     stack: FiberStack,
     inner: imp::Fiber,
diff --git a/crates/fiber/src/unix.rs b/crates/fiber/src/unix.rs
index b148ff324..7b6534344 100644
--- a/crates/fiber/src/unix.rs
+++ b/crates/fiber/src/unix.rs
@@ -31,21 +31,22 @@
 
 #![allow(unused_macros)]
 
-use crate::RunResult;
+use crate::{RunResult, RuntimeFiberStack};
 use std::cell::Cell;
 use std::io;
 use std::ops::Range;
 use std::ptr;
 
-#[derive(Debug)]
-pub struct FiberStack {
-    // The top of the stack; for stacks allocated by the fiber implementation itself,
-    // the base address of the allocation will be `top.sub(len.unwrap())`
-    top: *mut u8,
-    // The length of the stack
-    len: usize,
-    // whether or not this stack was mmap'd
-    mmap: bool,
+pub enum FiberStack {
+    Default {
+        // The top of the stack; for stacks allocated by the fiber implementation itself,
+        // the base address of the allocation will be `top.sub(len.unwrap())`
+        top: *mut u8,
+        // The length of the stack
+        len: usize,
+        mmap: bool,
+    },
+    Custom(Box<dyn RuntimeFiberStack>),
 }
 
 impl FiberStack {
@@ -75,7 +76,7 @@ impl FiberStack {
                 rustix::mm::MprotectFlags::READ | rustix::mm::MprotectFlags::WRITE,
             )?;
 
-            Ok(Self {
+            Ok(Self::Default {
                 top: mmap.cast::<u8>().add(mmap_len),
                 len: mmap_len,
                 mmap: true,
@@ -84,28 +85,77 @@ impl FiberStack {
     }
 
     pub unsafe fn from_raw_parts(base: *mut u8, len: usize) -> io::Result<Self> {
-        Ok(Self {
+        Ok(Self::Default {
             top: base.add(len),
             len,
             mmap: false,
         })
     }
 
+    pub fn from_custom(custom: Box<dyn RuntimeFiberStack>) -> io::Result<Self> {
+        Ok(Self::Custom(custom))
+    }
+
     pub fn top(&self) -> Option<*mut u8> {
-        Some(self.top)
+        Some(match self {
+            FiberStack::Default {
+                top,
+                len: _,
+                mmap: _,
+            } => *top,
+            FiberStack::Custom(r) => {
+                let top = r.top();
+                let page_size = rustix::param::page_size();
+                assert!(
+                    top.align_offset(page_size) == 0,
+                    "expected fiber stack top ({}) to be page aligned ({})",
+                    top as usize,
+                    page_size
+                );
+                top
+            }
+        })
     }
 
     pub fn range(&self) -> Option<Range<usize>> {
-        let base = unsafe { self.top.sub(self.len) as usize };
-        Some(base..base + self.len)
+        Some(match self {
+            FiberStack::Default { top, len, mmap: _ } => {
+                let base = unsafe { top.sub(*len) as usize };
+                base..base + len
+            }
+            FiberStack::Custom(s) => {
+                let range = s.range();
+                let page_size = rustix::param::page_size();
+                let start_ptr = range.start as *const u8;
+                assert!(
+                    start_ptr.align_offset(page_size) == 0,
+                    "expected fiber stack end ({}) to be page aligned ({})",
+                    range.start as usize,
+                    page_size
+                );
+                let end_ptr = range.end as *const u8;
+                assert!(
+                    end_ptr.align_offset(page_size) == 0,
+                    "expected fiber stack start ({}) to be page aligned ({})",
+                    range.end as usize,
+                    page_size
+                );
+                range
+            }
+        })
     }
 }
 
 impl Drop for FiberStack {
     fn drop(&mut self) {
         unsafe {
-            if self.mmap {
-                let ret = rustix::mm::munmap(self.top.sub(self.len) as _, self.len);
+            if let FiberStack::Default {
+                top,
+                len,
+                mmap: true,
+            } = self
+            {
+                let ret = rustix::mm::munmap(top.sub(*len) as _, *len);
                 debug_assert!(ret.is_ok());
             }
         }
@@ -148,7 +198,7 @@ impl Fiber {
     {
         unsafe {
             let data = Box::into_raw(Box::new(func)).cast();
-            wasmtime_fiber_init(stack.top, fiber_start::<F, A, B, C>, data);
+            wasmtime_fiber_init(stack.top().unwrap(), fiber_start::<F, A, B, C>, data);
         }
 
         Ok(Self)
@@ -160,10 +210,10 @@ impl Fiber {
             // stack, otherwise known as our reserved slot for this information.
             //
             // In the diagram above this is updating address 0xAff8
-            let addr = stack.top.cast::<usize>().offset(-1);
+            let addr = stack.top().unwrap().cast::<usize>().offset(-1);
             addr.write(result as *const _ as usize);
 
-            wasmtime_fiber_switch(stack.top);
+            wasmtime_fiber_switch(stack.top().unwrap());
 
             // null this out to help catch use-after-free
             addr.write(0);
diff --git a/crates/fiber/src/windows.rs b/crates/fiber/src/windows.rs
index 2e09f28b7..d04172704 100644
--- a/crates/fiber/src/windows.rs
+++ b/crates/fiber/src/windows.rs
@@ -1,4 +1,4 @@
-use crate::RunResult;
+use crate::{RunResult, RuntimeFiberStack};
 use std::cell::Cell;
 use std::ffi::c_void;
 use std::io;
@@ -19,6 +19,10 @@ impl FiberStack {
         Err(io::Error::from_raw_os_error(ERROR_NOT_SUPPORTED as i32))
     }
 
+    pub fn from_custom(_custom: Box<dyn RuntimeFiberStack>) -> io::Result<Self> {
+        Err(io::Error::from_raw_os_error(ERROR_NOT_SUPPORTED as i32))
+    }
+
     pub fn top(&self) -> Option<*mut u8> {
         None
     }
diff --git a/crates/runtime/src/instance/allocator/on_demand.rs b/crates/runtime/src/instance/allocator/on_demand.rs
index 935074729..5c7965304 100644
--- a/crates/runtime/src/instance/allocator/on_demand.rs
+++ b/crates/runtime/src/instance/allocator/on_demand.rs
@@ -12,6 +12,9 @@ use wasmtime_environ::{
     DefinedMemoryIndex, DefinedTableIndex, HostPtr, MemoryPlan, Module, TablePlan, VMOffsets,
 };
 
+#[cfg(feature = "async")]
+use wasmtime_fiber::RuntimeFiberStackCreator;
+
 #[cfg(feature = "component-model")]
 use wasmtime_environ::{
     component::{Component, VMComponentOffsets},
@@ -23,19 +26,29 @@ use wasmtime_environ::{
 pub struct OnDemandInstanceAllocator {
     mem_creator: Option<Arc<dyn RuntimeMemoryCreator>>,
     #[cfg(feature = "async")]
+    stack_creator: Option<Arc<dyn RuntimeFiberStackCreator>>,
+    #[cfg(feature = "async")]
     stack_size: usize,
 }
 
 impl OnDemandInstanceAllocator {
     /// Creates a new on-demand instance allocator.
     pub fn new(mem_creator: Option<Arc<dyn RuntimeMemoryCreator>>, stack_size: usize) -> Self {
-        let _ = stack_size; // suppress unused warnings w/o async feature
+        let _ = stack_size; // suppress warnings when async feature is disabled.
         Self {
             mem_creator,
             #[cfg(feature = "async")]
+            stack_creator: None,
+            #[cfg(feature = "async")]
             stack_size,
         }
     }
+
+    /// Set the stack creator.
+    #[cfg(feature = "async")]
+    pub fn set_stack_creator(&mut self, stack_creator: Arc<dyn RuntimeFiberStackCreator>) {
+        self.stack_creator = Some(stack_creator);
+    }
 }
 
 impl Default for OnDemandInstanceAllocator {
@@ -43,6 +56,8 @@ impl Default for OnDemandInstanceAllocator {
         Self {
             mem_creator: None,
             #[cfg(feature = "async")]
+            stack_creator: None,
+            #[cfg(feature = "async")]
             stack_size: 0,
         }
     }
@@ -141,8 +156,13 @@ unsafe impl InstanceAllocatorImpl for OnDemandInstanceAllocator {
         if self.stack_size == 0 {
             anyhow::bail!("fiber stacks are not supported by the allocator")
         }
-
-        let stack = wasmtime_fiber::FiberStack::new(self.stack_size)?;
+        let stack = match &self.stack_creator {
+            Some(stack_creator) => {
+                let stack = stack_creator.new_stack(self.stack_size)?;
+                wasmtime_fiber::FiberStack::from_custom(stack)
+            }
+            None => wasmtime_fiber::FiberStack::new(self.stack_size),
+        }?;
         Ok(stack)
     }
 
diff --git a/crates/runtime/src/instance/allocator/pooling/stack_pool.rs b/crates/runtime/src/instance/allocator/pooling/stack_pool.rs
index 8265cfcf9..edb169f34 100644
--- a/crates/runtime/src/instance/allocator/pooling/stack_pool.rs
+++ b/crates/runtime/src/instance/allocator/pooling/stack_pool.rs
@@ -213,7 +213,7 @@ mod tests {
 
         assert_eq!(pool.index_allocator.testing_freelist(), []);
 
-        pool.allocate().unwrap_err();
+        assert!(pool.allocate().is_err(), "allocation should fail");
 
         for stack in stacks {
             unsafe {
diff --git a/crates/runtime/src/memory.rs b/crates/runtime/src/memory.rs
index 84d568859..59cac8cab 100644
--- a/crates/runtime/src/memory.rs
+++ b/crates/runtime/src/memory.rs
@@ -565,6 +565,7 @@ impl SharedMemory {
     /// Implementation of `memory.atomic.notify` for this shared memory.
     pub fn atomic_notify(&self, addr_index: u64, count: u32) -> Result<u32, Trap> {
         validate_atomic_addr(&self.0.def.0, addr_index, 4, 4)?;
+        log::trace!("memory.atomic.notify(addr={addr_index}, count={count})");
         Ok(self.0.spot.unpark(addr_index, count))
     }
 
@@ -576,13 +577,18 @@ impl SharedMemory {
         timeout: Option<Instant>,
     ) -> Result<WaitResult, Trap> {
         let addr = validate_atomic_addr(&self.0.def.0, addr_index, 4, 4)?;
+        log::trace!(
+            "memory.atomic.wait32(addr={addr_index}, expected={expected}, timeout={timeout:?})"
+        );
+
         // SAFETY: `addr_index` was validated by `validate_atomic_addr` above.
         assert!(std::mem::size_of::<AtomicU32>() == 4);
         assert!(std::mem::align_of::<AtomicU32>() <= 4);
         let atomic = unsafe { &*(addr as *const AtomicU32) };
 
-        // We want the sequential consistency of `SeqCst` to ensure that the `load` sees the value that the `notify` will/would see.
-        // All WASM atomic operations are also `SeqCst`.
+        // We want the sequential consistency of `SeqCst` to ensure that the
+        // `load` sees the value that the `notify` will/would see. All WASM
+        // atomic operations are also `SeqCst`.
         let validate = || atomic.load(Ordering::SeqCst) == expected;
 
         Ok(self.0.spot.park(addr_index, validate, timeout))
@@ -596,6 +602,10 @@ impl SharedMemory {
         timeout: Option<Instant>,
     ) -> Result<WaitResult, Trap> {
         let addr = validate_atomic_addr(&self.0.def.0, addr_index, 8, 8)?;
+        log::trace!(
+            "memory.atomic.wait64(addr={addr_index}, expected={expected}, timeout={timeout:?})"
+        );
+
         // SAFETY: `addr_index` was validated by `validate_atomic_addr` above.
         assert!(std::mem::size_of::<AtomicU64>() == 8);
         assert!(std::mem::align_of::<AtomicU64>() <= 8);
diff --git a/crates/wasmtime/src/config.rs b/crates/wasmtime/src/config.rs
index 1fdec8937..3b9aa6cad 100644
--- a/crates/wasmtime/src/config.rs
+++ b/crates/wasmtime/src/config.rs
@@ -16,6 +16,11 @@ use wasmtime_environ::Tunables;
 use wasmtime_jit::profiling::{self, ProfilingAgent};
 use wasmtime_runtime::{mpk, InstanceAllocator, OnDemandInstanceAllocator, RuntimeMemoryCreator};
 
+#[cfg(feature = "async")]
+use crate::stack::{StackCreator, StackCreatorProxy};
+#[cfg(feature = "async")]
+use wasmtime_fiber::RuntimeFiberStackCreator;
+
 pub use wasmtime_environ::CacheStore;
 pub use wasmtime_runtime::MpkEnabled;
 
@@ -105,6 +110,8 @@ pub struct Config {
     pub(crate) native_unwind_info: Option<bool>,
     #[cfg(feature = "async")]
     pub(crate) async_stack_size: usize,
+    #[cfg(feature = "async")]
+    pub(crate) stack_creator: Option<Arc<dyn RuntimeFiberStackCreator>>,
     pub(crate) async_support: bool,
     pub(crate) module_version: ModuleVersionStrategy,
     pub(crate) parallel_compilation: bool,
@@ -199,6 +206,8 @@ impl Config {
             features: WasmFeatures::default(),
             #[cfg(feature = "async")]
             async_stack_size: 2 << 20,
+            #[cfg(feature = "async")]
+            stack_creator: None,
             async_support: false,
             module_version: ModuleVersionStrategy::default(),
             parallel_compilation: !cfg!(miri),
@@ -1084,6 +1093,17 @@ impl Config {
         self
     }
 
+    /// Sets a custom stack creator.
+    ///
+    /// Custom memory creators are used when creating creating async instance stacks for
+    /// the on-demand instance allocation strategy.
+    #[cfg(feature = "async")]
+    #[cfg_attr(nightlydoc, doc(cfg(feature = "async")))]
+    pub fn with_host_stack(&mut self, stack_creator: Arc<dyn StackCreator>) -> &mut Self {
+        self.stack_creator = Some(Arc::new(StackCreatorProxy(stack_creator)));
+        self
+    }
+
     /// Sets the instance allocation strategy to use.
     ///
     /// When using the pooling instance allocation strategy, all linear memories
@@ -1573,10 +1593,18 @@ impl Config {
         let stack_size = 0;
 
         match &self.allocation_strategy {
-            InstanceAllocationStrategy::OnDemand => Ok(Box::new(OnDemandInstanceAllocator::new(
-                self.mem_creator.clone(),
-                stack_size,
-            ))),
+            InstanceAllocationStrategy::OnDemand => {
+                #[allow(unused_mut)]
+                let mut allocator = Box::new(OnDemandInstanceAllocator::new(
+                    self.mem_creator.clone(),
+                    stack_size,
+                ));
+                #[cfg(feature = "async")]
+                if let Some(stack_creator) = &self.stack_creator {
+                    allocator.set_stack_creator(stack_creator.clone());
+                }
+                Ok(allocator)
+            }
             #[cfg(feature = "pooling-allocator")]
             InstanceAllocationStrategy::Pooling(config) => {
                 let mut config = config.config;
diff --git a/crates/wasmtime/src/lib.rs b/crates/wasmtime/src/lib.rs
index 1abfd28c0..f59ea2bcf 100644
--- a/crates/wasmtime/src/lib.rs
+++ b/crates/wasmtime/src/lib.rs
@@ -414,6 +414,9 @@ mod types;
 mod v128;
 mod values;
 
+#[cfg(feature = "async")]
+mod stack;
+
 pub use crate::config::*;
 pub use crate::coredump::*;
 pub use crate::engine::*;
@@ -437,6 +440,9 @@ pub use crate::types::*;
 pub use crate::v128::V128;
 pub use crate::values::*;
 
+#[cfg(feature = "async")]
+pub use crate::stack::*;
+
 /// A convenience wrapper for `Result<T, anyhow::Error>`.
 ///
 /// This type can be used to interact with `wasmtimes`'s extensive use
diff --git a/crates/wasmtime/src/stack.rs b/crates/wasmtime/src/stack.rs
new file mode 100644
index 000000000..fde7dd756
--- /dev/null
+++ b/crates/wasmtime/src/stack.rs
@@ -0,0 +1,73 @@
+use std::{ops::Range, sync::Arc};
+
+use anyhow::Error;
+use wasmtime_fiber::{RuntimeFiberStack, RuntimeFiberStackCreator};
+
+/// A stack creator. Can be used to provide a stack creator to wasmtime
+/// which supplies stacks for async support.
+///
+/// # Safety
+///
+/// This trait is unsafe, as memory safety depends on a proper implementation
+/// of memory management. Stacks created by the StackCreator should always be
+/// treated as owned by an wasmtime instance, and any modification of them
+/// outside of wasmtime invoked routines is unsafe and may lead to corruption.
+///
+/// Note that this is a relatively new and experimental feature and it is
+/// recommended to be familiar with wasmtime runtime code to use it.
+pub unsafe trait StackCreator: Send + Sync {
+    /// Create a new `StackMemory` object with the specified size.
+    ///
+    /// The `size` parameter is the expected size of the stack without any guard pages.
+    ///
+    /// Note there should be at least one guard page of protected memory at the bottom
+    /// of the stack to catch potential stack overflow scenarios. Additionally, stacks should be
+    /// page aligned.
+    fn new_stack(&self, size: usize) -> Result<Box<dyn StackMemory>, Error>;
+}
+
+#[derive(Clone)]
+pub(crate) struct StackCreatorProxy(pub Arc<dyn StackCreator>);
+
+unsafe impl RuntimeFiberStackCreator for StackCreatorProxy {
+    fn new_stack(&self, size: usize) -> Result<Box<dyn RuntimeFiberStack>, Error> {
+        let stack = self.0.new_stack(size)?;
+        Ok(Box::new(FiberStackProxy(stack)) as Box<dyn RuntimeFiberStack>)
+    }
+}
+
+/// A stack memory. This trait provides an interface for raw memory buffers
+/// which are used by wasmtime inside of stacks which wasmtime executes
+/// WebAssembly in for async support. By implementing this trait together
+/// with StackCreator, one can supply wasmtime with custom allocated host
+/// managed stacks.
+///
+/// # Safety
+///
+/// The memory should be page aligned and a multiple of page size.
+/// To prevent possible silent overflows, the memory should be protected by a
+/// guard page. Additionally the safety concerns explained in ['Memory'], for
+/// accessing the memory apply here as well.
+///
+/// Note that this is a relatively new and experimental feature and it is
+/// recommended to be familiar with wasmtime runtime code to use it.
+pub unsafe trait StackMemory: Send + Sync {
+    /// The top of the allocated stack.
+    ///
+    /// This address should be page aligned.
+    fn top(&self) -> *mut u8;
+    /// The range of where this stack resides in memory, excluding guard pages.
+    fn range(&self) -> Range<usize>;
+}
+
+pub(crate) struct FiberStackProxy(pub Box<dyn StackMemory>);
+
+unsafe impl RuntimeFiberStack for FiberStackProxy {
+    fn top(&self) -> *mut u8 {
+        self.0.top()
+    }
+
+    fn range(&self) -> Range<usize> {
+        self.0.range()
+    }
+}
diff --git a/crates/wasmtime/src/trampoline.rs b/crates/wasmtime/src/trampoline.rs
index 7f2499fb3..2c28509c8 100644
--- a/crates/wasmtime/src/trampoline.rs
+++ b/crates/wasmtime/src/trampoline.rs
@@ -41,15 +41,15 @@ fn create_handle(
         let module = Arc::new(module);
         let runtime_info =
             &BareModuleInfo::maybe_imported_func(module, one_signature).into_traitobj();
-        let handle = OnDemandInstanceAllocator::new(config.mem_creator.clone(), 0)
-            .allocate_module(InstanceAllocationRequest {
-                imports,
-                host_state,
-                store: StorePtr::new(store.traitobj()),
-                runtime_info,
-                wmemcheck: false,
-                pkey: None,
-            })?;
+        let allocator = OnDemandInstanceAllocator::new(config.mem_creator.clone(), 0);
+        let handle = allocator.allocate_module(InstanceAllocationRequest {
+            imports,
+            host_state,
+            store: StorePtr::new(store.traitobj()),
+            runtime_info,
+            wmemcheck: false,
+            pkey: None,
+        })?;
 
         Ok(store.add_dummy_instance(handle))
     }
diff --git a/docs/WASI-intro.md b/docs/WASI-intro.md
index 187f3003c..8e2507dbd 100644
--- a/docs/WASI-intro.md
+++ b/docs/WASI-intro.md
@@ -40,7 +40,7 @@ Currently the options are [Wasmtime] and the [browser polyfill], though we
 intend WASI to be implementable in many wasm VMs.
 
 [Wasmtime]: https://github.com/bytecodealliance/wasmtime
-[browser polyfill]: https://wasi.dev/polyfill/
+[browser polyfill]: https://github.com/bjorn3/browser_wasi_shim
 
 ### Wasmtime
 
@@ -52,9 +52,7 @@ or `cargo run --bin wasmtime foo.wasm`.
 
 ### The browser polyfill
 
-The polyfill is online [here](https://wasi.dev/polyfill/).
-
-The source is [here](https://github.com/bytecodealliance/wasmtime/tree/main/crates/wasi-c/js-polyfill).
+The source is [here](https://github.com/bjorn3/browser_wasi_shim).
 
 ## Where can I learn more?
 


Test patch:
diff --git a/cranelift/filetests/filetests/isa/riscv64/zca.clif b/cranelift/filetests/filetests/isa/riscv64/zca.clif
index 4726ec5d4..b49fd2d63 100644
--- a/cranelift/filetests/filetests/isa/riscv64/zca.clif
+++ b/cranelift/filetests/filetests/isa/riscv64/zca.clif
@@ -432,12 +432,16 @@ block0(v0: i64):
 ;   c.srli a0, 0x14
 ;   c.jr ra
 
-function %c_addi4spn() -> i64 {
-  ss0 = explicit_slot 64
+function %c_addi4spn() -> i64, i64, i64, i64, i64 {
+  ss0 = explicit_slot 2048
 
 block0:
-  v0 = stack_addr.i64 ss0+24
-  return v0
+  v0 = stack_addr.i64 ss0+1020
+  v1 = stack_addr.i64 ss0+4
+  v2 = stack_addr.i64 ss0+512
+  v4 = stack_addr.i64 ss0+256
+  v5 = stack_addr.i64 ss0+64
+  return v0, v1, v2, v4, v5
 }
 
 ; VCode:
@@ -445,10 +449,69 @@ block0:
 ;   sd ra,8(sp)
 ;   sd fp,0(sp)
 ;   mv fp,sp
-;   add sp,-64
+;   add sp,-2048
+; block0:
+;   load_addr a1,1020(nominal_sp)
+;   mv a5,a1
+;   load_addr a1,4(nominal_sp)
+;   load_addr a2,512(nominal_sp)
+;   load_addr a3,256(nominal_sp)
+;   load_addr a4,64(nominal_sp)
+;   sd a2,0(a0)
+;   sd a3,8(a0)
+;   sd a4,16(a0)
+;   mv a0,a5
+;   add sp,+2048
+;   ld ra,8(sp)
+;   ld fp,0(sp)
+;   add sp,+16
+;   ret
+;
+; Disassembled:
+; block0: ; offset 0x0
+;   c.addi16sp sp, -0x10
+;   c.sdsp ra, 8(sp)
+;   c.sdsp s0, 0(sp)
+;   c.mv s0, sp
+;   addi sp, sp, -0x800
+; block1: ; offset 0xc
+;   c.addi4spn a1, sp, 0x3fc
+;   c.mv a5, a1
+;   c.addi4spn a1, sp, 4
+;   c.addi4spn a2, sp, 0x200
+;   c.addi4spn a3, sp, 0x100
+;   c.addi4spn a4, sp, 0x40
+;   c.sd a2, 0(a0)
+;   c.sd a3, 8(a0)
+;   c.sd a4, 0x10(a0)
+;   c.mv a0, a5
+;   c.lui t6, 1
+;   addi t6, t6, -0x800
+;   add sp, t6, sp
+;   c.ldsp ra, 8(sp)
+;   c.ldsp s0, 0(sp)
+;   c.addi16sp sp, 0x10
+;   c.jr ra
+
+function %out_of_bounds_c_addi4spn() -> i64, i64 {
+  ss0 = explicit_slot 2048
+
+block0:
+  v0 = stack_addr.i64 ss0+1024
+  v1 = stack_addr.i64 ss0+0
+  return v0, v1
+}
+
+; VCode:
+;   add sp,-16
+;   sd ra,8(sp)
+;   sd fp,0(sp)
+;   mv fp,sp
+;   add sp,-2048
 ; block0:
-;   load_addr a0,24(nominal_sp)
-;   add sp,+64
+;   load_addr a0,1024(nominal_sp)
+;   load_addr a1,0(nominal_sp)
+;   add sp,+2048
 ;   ld ra,8(sp)
 ;   ld fp,0(sp)
 ;   add sp,+16
@@ -460,10 +523,13 @@ block0:
 ;   c.sdsp ra, 8(sp)
 ;   c.sdsp s0, 0(sp)
 ;   c.mv s0, sp
-;   c.addi16sp sp, -0x40
-; block1: ; offset 0xa
-;   c.addi4spn a0, sp, 0x18
-;   c.addi16sp sp, 0x40
+;   addi sp, sp, -0x800
+; block1: ; offset 0xc
+;   addi a0, sp, 0x400
+;   c.mv a1, sp
+;   c.lui t6, 1
+;   addi t6, t6, -0x800
+;   add sp, t6, sp
 ;   c.ldsp ra, 8(sp)
 ;   c.ldsp s0, 0(sp)
 ;   c.addi16sp sp, 0x10
diff --git a/tests/all/main.rs b/tests/all/main.rs
index 54c8229b8..79ad24c60 100644
--- a/tests/all/main.rs
+++ b/tests/all/main.rs
@@ -30,6 +30,7 @@ mod module_serialize;
 mod name;
 mod pooling_allocator;
 mod relocs;
+mod stack_creator;
 mod stack_overflow;
 mod store;
 mod table;
diff --git a/tests/all/stack_creator.rs b/tests/all/stack_creator.rs
new file mode 100644
index 000000000..5b46d8942
--- /dev/null
+++ b/tests/all/stack_creator.rs
@@ -0,0 +1,155 @@
+#![cfg(all(not(target_os = "windows"), not(miri)))]
+use anyhow::{bail, Context};
+use std::{
+    alloc::{GlobalAlloc, Layout, System},
+    ops::Range,
+    ptr::NonNull,
+    sync::Arc,
+};
+use wasmtime::*;
+
+fn align_up(v: usize, align: usize) -> usize {
+    return (v + (align - 1)) & (!(align - 1));
+}
+
+struct CustomStack {
+    base: NonNull<u8>,
+    len: usize,
+}
+unsafe impl Send for CustomStack {}
+unsafe impl Sync for CustomStack {}
+impl CustomStack {
+    fn new(base: NonNull<u8>, len: usize) -> Self {
+        CustomStack { base, len }
+    }
+}
+unsafe impl StackMemory for CustomStack {
+    fn top(&self) -> *mut u8 {
+        unsafe { self.base.as_ptr().add(self.len) }
+    }
+    fn range(&self) -> Range<usize> {
+        let base = self.base.as_ptr() as usize;
+        base..base + self.len
+    }
+}
+
+// A creator that allocates stacks on the heap instead of mmap'ing.
+struct CustomStackCreator {
+    memory: NonNull<u8>,
+    size: usize,
+    layout: Layout,
+}
+
+unsafe impl Send for CustomStackCreator {}
+unsafe impl Sync for CustomStackCreator {}
+impl CustomStackCreator {
+    fn new() -> Result<Self> {
+        // 1MB
+        const MINIMUM_STACK_SIZE: usize = 1 * 1_024 * 1_024;
+        let page_size = rustix::param::page_size();
+        let size = align_up(MINIMUM_STACK_SIZE, page_size);
+        // Add an extra page for the guard page
+        let layout = Layout::from_size_align(size + page_size, page_size)
+            .context("unable to compute stack layout")?;
+        let memory = unsafe {
+            let mem = System.alloc(layout);
+            let notnull = NonNull::new(mem);
+            if let Some(mem) = notnull {
+                // Mark guard page as protected
+                rustix::mm::mprotect(
+                    mem.as_ptr().cast(),
+                    page_size,
+                    rustix::mm::MprotectFlags::empty(),
+                )?;
+            }
+            notnull
+        }
+        .context("unable to allocate stack memory")?;
+        Ok(CustomStackCreator {
+            memory,
+            size,
+            layout,
+        })
+    }
+    fn range(&self) -> Range<usize> {
+        let page_size = rustix::param::page_size();
+        let base = unsafe { self.memory.as_ptr().add(page_size) as usize };
+        base..base + self.size
+    }
+}
+impl Drop for CustomStackCreator {
+    fn drop(&mut self) {
+        let page_size = rustix::param::page_size();
+        unsafe {
+            // Unprotect the guard page as the allocator could reuse it.
+            rustix::mm::mprotect(
+                self.memory.as_ptr().cast(),
+                page_size,
+                rustix::mm::MprotectFlags::READ | rustix::mm::MprotectFlags::WRITE,
+            )
+            .unwrap();
+            System.dealloc(self.memory.as_ptr(), self.layout);
+        }
+    }
+}
+unsafe impl StackCreator for CustomStackCreator {
+    fn new_stack(&self, size: usize) -> Result<Box<dyn StackMemory>> {
+        if size != self.size {
+            bail!("must use the size we allocated for this stack memory creator");
+        }
+        let page_size = rustix::param::page_size();
+        // skip over the page size
+        let base_ptr = unsafe { self.memory.as_ptr().add(page_size) };
+        let base = NonNull::new(base_ptr).context("unable to compute stack base")?;
+        Ok(Box::new(CustomStack::new(base, self.size)))
+    }
+}
+
+fn config() -> (Store<()>, Arc<CustomStackCreator>) {
+    let stack_creator = Arc::new(CustomStackCreator::new().unwrap());
+    let mut config = Config::new();
+    config
+        .async_support(true)
+        .max_wasm_stack(stack_creator.size / 2)
+        .async_stack_size(stack_creator.size)
+        .with_host_stack(stack_creator.clone());
+    (
+        Store::new(&Engine::new(&config).unwrap(), ()),
+        stack_creator,
+    )
+}
+
+#[tokio::test]
+async fn called_on_custom_heap_stack() -> Result<()> {
+    let (mut store, stack_creator) = config();
+    let module = Module::new(
+        store.engine(),
+        r#"
+            (module
+                (import "host" "callback" (func $callback (result i64)))
+                (func $f (result i64) (call $callback))
+                (export "f" (func $f))
+            )
+        "#,
+    )?;
+
+    let ty = FuncType::new([], [ValType::I64]);
+    let host_func = Func::new(&mut store, ty, move |_caller, _params, results| {
+        let foo = 42;
+        // output an address on the stack
+        results[0] = Val::I64((&foo as *const i32) as usize as i64);
+        Ok(())
+    });
+    let export = wasmtime::Extern::Func(host_func);
+    let instance = Instance::new_async(&mut store, &module, &[export]).await?;
+    let mut results = [Val::I64(0)];
+    instance
+        .get_func(&mut store, "f")
+        .context("missing function export")?
+        .call_async(&mut store, &[], &mut results)
+        .await?;
+    // Make sure the stack address we wrote was within our custom stack range
+    let stack_address = results[0].i64().unwrap() as usize;
+    assert!(stack_creator.range().contains(&stack_address));
+    Ok(())
+}