Introduce API for custom stack memory

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

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,
+    }));
+}

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,

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 }
 

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,