Address review comments

Signed-off-by: Tyler Rockwood <rockwood@redpanda.com>

crates/fiber/src/unix.rs

@@ -38,19 +38,13 @@ use std::ops::Range;
 use std::ptr;
 
 pub enum FiberStack {
-    Mmap {
-        // 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,
-    },
-    Manual {
+    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>),
 }
@@ -82,17 +76,19 @@ impl FiberStack {
                 rustix::mm::MprotectFlags::READ | rustix::mm::MprotectFlags::WRITE,
             )?;
 
-            Ok(Self::Mmap {
+            Ok(Self::Default {
                 top: mmap.cast::<u8>().add(mmap_len),
                 len: mmap_len,
+                mmap: true,
             })
         }
     }
 
     pub unsafe fn from_raw_parts(base: *mut u8, len: usize) -> io::Result<Self> {
-        Ok(Self::Manual {
+        Ok(Self::Default {
             top: base.add(len),
             len,
+            mmap: false,
         })
     }
 
@@ -102,31 +98,63 @@ impl FiberStack {
 
     pub fn top(&self) -> Option<*mut u8> {
         Some(match self {
-            FiberStack::Mmap { top, len: _ } => *top,
-            FiberStack::Manual { top, len: _ } => *top,
-            FiberStack::Custom(r) => r.top(),
+            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>> {
         Some(match self {
-            FiberStack::Mmap { top, len } => {
+            FiberStack::Default { top, len, mmap: _ } => {
                 let base = unsafe { top.sub(*len) as usize };
                 base..base + len
             }
-            FiberStack::Manual { top, len } => {
-                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
             }
-            FiberStack::Custom(r) => r.range(),
         })
     }
 }
 
 impl Drop for FiberStack {
     fn drop(&mut self) {
         unsafe {
-            if let FiberStack::Mmap { top, len } = self {
+            if let FiberStack::Default {
+                top,
+                len,
+                mmap: true,
+            } = self
+            {
                 let ret = rustix::mm::munmap(top.sub(*len) as _, *len);
                 debug_assert!(ret.is_ok());
             }

crates/runtime/src/instance/allocator/on_demand.rs

@@ -33,22 +33,20 @@ pub struct OnDemandInstanceAllocator {
 
 impl OnDemandInstanceAllocator {
     /// Creates a new on-demand instance allocator.
-    #[cfg(feature = "async")]
-    pub fn new(
-        mem_creator: Option<Arc<dyn RuntimeMemoryCreator>>,
-        stack_creator: Option<Arc<dyn RuntimeFiberStackCreator>>,
-        stack_size: usize,
-    ) -> Self {
+    pub fn new(mem_creator: Option<Arc<dyn RuntimeMemoryCreator>>, stack_size: usize) -> Self {
         Self {
             mem_creator,
-            stack_creator,
+            #[cfg(feature = "async")]
+            stack_creator: None,
             #[cfg(feature = "async")]
             stack_size,
         }
     }
-    #[cfg(not(feature = "async"))]
-    pub fn new(mem_creator: Option<Arc<dyn RuntimeMemoryCreator>>) -> Self {
-        Self { mem_creator }
+
+    /// 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);
     }
 }
 

crates/wasmtime/src/config.rs

@@ -1587,16 +1587,18 @@ impl Config {
 
     pub(crate) fn build_allocator(&self) -> Result<Box<dyn InstanceAllocator + Send + Sync>> {
         match &self.allocation_strategy {
-            #[cfg(feature = "async")]
-            InstanceAllocationStrategy::OnDemand => Ok(Box::new(OnDemandInstanceAllocator::new(
-                self.mem_creator.clone(),
-                self.stack_creator.clone(),
-                self.async_stack_size,
-            ))),
-            #[cfg(not(feature = "async"))]
-            InstanceAllocationStrategy::OnDemand => Ok(Box::new(OnDemandInstanceAllocator::new(
-                self.mem_creator.clone(),
-            ))),
+            InstanceAllocationStrategy::OnDemand => {
+                #[allow(unused_mut)]
+                let mut allocator = Box::new(OnDemandInstanceAllocator::new(
+                    self.mem_creator.clone(),
+                    self.async_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) => {
                 #[cfg(feature = "async")]

crates/wasmtime/src/stack.rs

@@ -8,7 +8,7 @@ use wasmtime_fiber::{RuntimeFiberStack, RuntimeFiberStackCreator};
 ///
 /// # Safety
 ///
-/// This trait is unsafe, as memory safety depends on a proper implemenation
+/// 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.

crates/wasmtime/src/trampoline.rs

@@ -41,10 +41,7 @@ fn create_handle(
         let module = Arc::new(module);
         let runtime_info =
             &BareModuleInfo::maybe_imported_func(module, one_signature).into_traitobj();
-        #[cfg(feature = "async")]
-        let allocator = OnDemandInstanceAllocator::new(config.mem_creator.clone(), None, 0);
-        #[cfg(not(feature = "async"))]
-        let allocator = OnDemandInstanceAllocator::new(config.mem_creator.clone());
+        let allocator = OnDemandInstanceAllocator::new(config.mem_creator.clone(), 0);
         let handle = allocator.allocate_module(InstanceAllocationRequest {
             imports,
             host_state,