Move unix_stack_pool over to aligned byte counts (#9678)

This removes the last remaining use of `round_usize_up_to_host_pages`. I've
added a couple of extra asserts that I believe are justified.

crates/wasmtime/src/runtime/vm.rs

@@ -367,33 +367,6 @@ pub fn host_page_size() -> usize {
     };
 }
 
-/// Round the given byte size up to a multiple of the host OS page size.
-///
-/// Returns an error if rounding up overflows.
-///
-/// (Deprecated: consider switching to `HostAlignedByteCount`.)
-#[cfg(all(feature = "async", unix, not(miri)))]
-pub fn round_u64_up_to_host_pages(bytes: u64) -> Result<u64> {
-    let page_size = u64::try_from(crate::runtime::vm::host_page_size()).err2anyhow()?;
-    debug_assert!(page_size.is_power_of_two());
-    bytes
-        .checked_add(page_size - 1)
-        .ok_or_else(|| anyhow!(
-            "{bytes} is too large to be rounded up to a multiple of the host page size of {page_size}"
-        ))
-        .map(|val| val & !(page_size - 1))
-}
-
-/// Same as `round_u64_up_to_host_pages` but for `usize`s.
-///
-/// (Deprecated: consider switching to `HostAlignedByteCount`.)
-#[cfg(all(feature = "async", unix, not(miri)))]
-pub fn round_usize_up_to_host_pages(bytes: usize) -> Result<usize> {
-    let bytes = u64::try_from(bytes).err2anyhow()?;
-    let rounded = round_u64_up_to_host_pages(bytes)?;
-    Ok(usize::try_from(rounded).err2anyhow()?)
-}
-
 /// Result of `Memory::atomic_wait32` and `Memory::atomic_wait64`
 #[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub enum WaitResult {

crates/wasmtime/src/runtime/vm/instance/allocator/pooling/unix_stack_pool.rs

@@ -4,8 +4,7 @@ use super::index_allocator::{SimpleIndexAllocator, SlotId};
 use crate::prelude::*;
 use crate::runtime::vm::sys::vm::commit_pages;
 use crate::runtime::vm::{
-    mmap::AlignedLength, round_usize_up_to_host_pages, HostAlignedByteCount, Mmap,
-    PoolingInstanceAllocatorConfig,
+    mmap::AlignedLength, HostAlignedByteCount, Mmap, PoolingInstanceAllocatorConfig,
 };
 
 /// Represents a pool of execution stacks (used for the async fiber implementation).
@@ -26,7 +25,7 @@ pub struct StackPool {
     page_size: HostAlignedByteCount,
     index_allocator: SimpleIndexAllocator,
     async_stack_zeroing: bool,
-    async_stack_keep_resident: usize,
+    async_stack_keep_resident: HostAlignedByteCount,
 }
 
 impl StackPool {
@@ -80,7 +79,7 @@ impl StackPool {
             max_stacks,
             page_size,
             async_stack_zeroing: config.async_stack_zeroing,
-            async_stack_keep_resident: round_usize_up_to_host_pages(
+            async_stack_keep_resident: HostAlignedByteCount::new_rounded_up(
                 config.async_stack_keep_resident,
             )?,
             index_allocator: SimpleIndexAllocator::new(config.limits.total_stacks),
@@ -95,7 +94,7 @@ impl StackPool {
 
     /// Allocate a new fiber.
     pub fn allocate(&self) -> Result<wasmtime_fiber::FiberStack> {
-        if self.stack_size == 0 {
+        if self.stack_size.is_zero() {
             bail!("pooling allocator not configured to enable fiber stack allocation");
         }
 
@@ -109,20 +108,23 @@ impl StackPool {
 
         unsafe {
             // Remove the guard page from the size
-            let size_without_guard = self.stack_size.byte_count() - self.page_size.byte_count();
+            let size_without_guard = self.stack_size.checked_sub(self.page_size).expect(
+                "self.stack_size is host-page-aligned and is > 0,\
+                 so it must be >= self.page_size",
+            );
 
             let bottom_of_stack = self
                 .mapping
                 .as_ptr()
                 .add(self.stack_size.unchecked_mul(index).byte_count())
                 .cast_mut();
 
-            commit_pages(bottom_of_stack, size_without_guard)?;
+            commit_pages(bottom_of_stack, size_without_guard.byte_count())?;
 
             let stack = wasmtime_fiber::FiberStack::from_raw_parts(
                 bottom_of_stack,
                 self.page_size.byte_count(),
-                size_without_guard,
+                size_without_guard.byte_count(),
             )?;
             Ok(stack)
         }
@@ -134,6 +136,11 @@ impl StackPool {
     /// that should be decommitted. It is the caller's responsibility to ensure
     /// that those decommits happen before this stack is reused.
     ///
+    /// # Panics
+    ///
+    /// `zero_stack` panics if the passed in `stack` was not created by
+    /// [`Self::allocate`].
+    ///
     /// # Safety
     ///
     /// The stack must no longer be in use, and ready for returning to the pool
@@ -144,6 +151,11 @@ impl StackPool {
         mut decommit: impl FnMut(*mut u8, usize),
     ) {
         assert!(stack.is_from_raw_parts());
+        assert!(
+            !self.stack_size.is_zero(),
+            "pooling allocator not configured to enable fiber stack allocation \
+             (Self::allocate should have returned an error)"
+        );
 
         if !self.async_stack_zeroing {
             return;
@@ -160,9 +172,12 @@ impl StackPool {
             "fiber stack top pointer not in range"
         );
 
-        // Remove the guard page from the size
-        let stack_size = self.stack_size.byte_count() - self.page_size.byte_count();
-        let bottom_of_stack = top - stack_size;
+        // Remove the guard page from the size.
+        let stack_size = self.stack_size.checked_sub(self.page_size).expect(
+            "self.stack_size is host-page-aligned and is > 0,\
+             so it must be >= self.page_size",
+        );
+        let bottom_of_stack = top - stack_size.byte_count();
         let start_of_stack = bottom_of_stack - self.page_size.byte_count();
         assert!(start_of_stack >= base && start_of_stack < (base + len));
         assert!((start_of_stack - base) % self.stack_size.byte_count() == 0);
@@ -175,14 +190,17 @@ impl StackPool {
         // * madvise for the whole range incurs expensive future page faults
         // * most threads probably don't use most of the stack anyway
         let size_to_memset = stack_size.min(self.async_stack_keep_resident);
+        let rest = stack_size
+            .checked_sub(size_to_memset)
+            .expect("stack_size >= size_to_memset");
         std::ptr::write_bytes(
-            (bottom_of_stack + stack_size - size_to_memset) as *mut u8,
+            (bottom_of_stack + rest.byte_count()) as *mut u8,
             0,
-            size_to_memset,
+            size_to_memset.byte_count(),
         );
 
         // Use the system to reset remaining stack pages to zero.
-        decommit(bottom_of_stack as _, stack_size - size_to_memset);
+        decommit(bottom_of_stack as _, rest.byte_count());
     }
 
     /// Deallocate a previously-allocated fiber.