Basic stack pooling allocator (bytecodealliance#169)

This patch adds a basic stack pooling allocator for WasmFX stacks. It is
gated behind a feature flag `wasmfx_pooling_allocator`, which is
disabled by default. By default
stacks are safe, meaning they have a page guard attached. The old unsafe
stacks are gated behind the feature flag `unsafe_wasmfx_stacks`.

In addition, this patch also adds missing WasmFX features to
`wasmtime-c-api` and `wasmtime-c-api-impl`.

Resolves bytecodealliance#91. Resolves bytecodealliance#26.

.github/workflows/main.yml

@@ -421,6 +421,11 @@ jobs:
       env:
         CARGO_PROFILE_DEV_DEBUG_ASSERTIONS: false
 
+    # Check a few WasmFX configurations
+    - run: cargo check -p wasmtime --no-default-features --features wasmfx_pooling_allocator,runtime,gc,component-model
+    - run: cargo check --features wasmfx_pooling_allocator
+    - run: cargo check --features wasmfx_baseline,wasmfx_pooling_allocator
+
     # Check whether `wasmtime` cross-compiles to x86_64-unknown-freebsd
     # TODO: We aren't building with default features since the `ittapi` crate fails to compile on freebsd.
     - run: rustup target add x86_64-unknown-freebsd
@@ -748,6 +753,17 @@ jobs:
       env:
         RUST_BACKTRACE: 1
 
+    # Test WasmFX allocators
+    - name: WasmFX pooling allocator
+      run: cargo test --features=default,wasmfx_pooling_allocator
+      env:
+        RUST_BACKTRACE: 1
+
+    - name: WasmFX pooling allocator (baseline)
+      run: cargo test --features=default,wasmfx_baseline,wasmfx_pooling_allocator
+      env:
+        RUST_BACKTRACE: 1
+
     # common logic to cancel the entire run if this job fails
     - run: gh run cancel ${{ github.run_id }}
       if: failure() && github.event_name != 'pull_request'

Cargo.toml

@@ -384,7 +384,20 @@ unsafe_disable_continuation_linearity_check = [
 # The following toggles the baseline implementation of typed continuations.
 wasmfx_baseline = [
   "wasmtime-cranelift/wasmfx_baseline",
-  "wasmtime/wasmfx_baseline"
+  "wasmtime/wasmfx_baseline",
+]
+
+# Toggle the WasmFX pooling allocator
+wasmfx_pooling_allocator = [
+  "wasmtime/wasmfx_pooling_allocator",
+  "pooling-allocator",
+]
+
+# Enable the old-style unsafe malloc'd stacks (has no effect for
+# typed_continuations_baseline_implementation or when
+# wasmfx_pooling_allocator is toggled)
+unsafe_wasmfx_stacks = [
+  "wasmtime/unsafe_wasmfx_stacks",
 ]
 
 # ========================================

crates/c-api/Cargo.toml

@@ -53,6 +53,27 @@ threads = ["wasmtime/threads"]
 gc = ["wasmtime/gc"]
 cranelift = ['wasmtime/cranelift']
 winch = ['wasmtime/winch']
+
+# Enable unsafe mutable continuations (temporary hack to workaround
+# the garbage generated by continuation objects).
+unsafe_disable_continuation_linearity_check = ["wasmtime/unsafe_disable_continuation_linearity_check"]
+
+# Toggle the baseline implementation of WasmFX
+wasmfx_baseline = [
+  "wasmtime/wasmfx_baseline"
+]
+
+# Toggle the WasmFX stack pooling allocator
+wasmfx_pooling_allocator = [
+  "wasmtime/wasmfx_pooling_allocator",
+  "wasmtime/pooling-allocator"
+]
+
+# Enable the old-style unsafe malloc'd stacks
+unsafe_wasmfx_stacks = ["wasmtime/unsafe_wasmfx_stacks"]
+
+
 # ... if you add a line above this be sure to also change:
 #
 #   crates/c-api/artifact/Cargo.toml
+

crates/c-api/artifact/Cargo.toml

@@ -54,5 +54,22 @@ threads = ["wasmtime-c-api/threads"]
 gc = ["wasmtime-c-api/gc"]
 cranelift = ["wasmtime-c-api/cranelift"]
 winch = ["wasmtime-c-api/winch"]
+
+# Enable unsafe mutable continuations (temporary hack to workaround
+# the garbage generated by continuation objects).
+unsafe_disable_continuation_linearity_check = [
+  "wasmtime-c-api/unsafe_disable_continuation_linearity_check"
+]
+
+# Toggle the baseline implementation of WasmFX
+wasmfx_baseline = ["wasmtime-c-api/wasmfx_baseline"]
+
+# Toggle the WasmFX stack pooling allocator
+wasmfx_pooling_allocator = ["wasmtime-c-api/wasmfx_pooling_allocator"]
+
+# Enable the old-style unsafe malloc'd stacks
+unsafe_wasmfx_stacks = ["wasmtime-c-api/unsafe_wasmfx_stacks"]
+
+
 # ... if you add a line above this be sure to read the comment at the end of
 # `default`

crates/wasmtime/Cargo.toml

@@ -267,7 +267,21 @@ std = [
 
 # Enable unsafe mutable continuations (temporary hack to workaround
 # the garbage generated by continuation objects).
-unsafe_disable_continuation_linearity_check = []
+unsafe_disable_continuation_linearity_check = [
+  "wasmtime-cranelift?/unsafe_disable_continuation_linearity_check"
+]
 
 # Toggle the baseline implementation of WasmFX
-wasmfx_baseline = ["async", "wasmtime-cranelift?/wasmfx_baseline"]
+wasmfx_baseline = [
+  "async",
+  "wasmtime-cranelift?/wasmfx_baseline"
+]
+
+# Toggle the WasmFX stack pooling allocator
+wasmfx_pooling_allocator = [
+  "pooling-allocator",
+]
+
+# Enable the old-style unsafe malloc'd stacks
+unsafe_wasmfx_stacks = []
+

crates/wasmtime/src/runtime/vm/continuation.rs

@@ -20,7 +20,6 @@ pub mod optimized {
     use super::stack_chain::StackChain;
     use crate::runtime::vm::{
         fibre::Fiber,
-        fibre::FiberStack,
         vmcontext::{VMFuncRef, ValRaw},
         Instance, TrapReason,
     };
@@ -32,6 +31,7 @@ pub mod optimized {
 
     /// Fibers used for continuations
     pub type ContinuationFiber = Fiber;
+    pub type FiberStack = crate::runtime::vm::fibre::FiberStack;
 
     /// TODO
     #[repr(C)]
@@ -90,11 +90,12 @@ pub mod optimized {
 
     /// TODO
     #[inline(always)]
-    pub fn drop_cont_ref(_instance: &mut Instance, contref: *mut VMContRef) {
+    pub fn drop_cont_ref(instance: &mut Instance, contref: *mut VMContRef) {
         // Note that continuation references do not own their parents, hence we ignore
         // parent fields here.
 
         let contref: Box<VMContRef> = unsafe { Box::from_raw(contref) };
+        instance.wasmfx_deallocate_stack(contref.fiber.stack());
         if contref.args.data.is_null() {
             debug_assert!(contref.args.length as usize == 0);
             debug_assert!(contref.args.capacity as usize == 0);
@@ -142,7 +143,7 @@ pub mod optimized {
         let red_zone_size = wasmfx_config.red_zone_size;
 
         let fiber = {
-            let stack = FiberStack::malloc(stack_size).map_err(|_error| {
+            let stack = instance.wasmfx_allocate_stack().map_err(|_error| {
                 TrapReason::user_without_backtrace(anyhow::anyhow!(
                     "Fiber stack allocation failed!"
                 ))
@@ -302,26 +303,23 @@ pub mod optimized {
         assert_eq!(memoffset::offset_of!(VMContRef, state), vm_cont_ref::STATE);
 
         assert_eq!(
-            core::mem::size_of::<ContinuationFiber>(),
+            std::mem::size_of::<ContinuationFiber>(),
             CONTINUATION_FIBER_SIZE
         );
-        assert_eq!(core::mem::size_of::<StackChain>(), STACK_CHAIN_SIZE);
+        assert_eq!(std::mem::size_of::<StackChain>(), STACK_CHAIN_SIZE);
     }
 }
 
-//
-// Baseline implementation
-//
 #[cfg(feature = "wasmfx_baseline")]
 pub mod baseline {
     use super::stack_chain::{StackChain, StackLimits};
     use crate::runtime::vm::{Instance, TrapReason, VMFuncRef, VMOpaqueContext, ValRaw};
     use core::{cell::Cell, cell::RefCell, cmp, mem};
-    use wasmtime_continuations::DEFAULT_FIBER_SIZE;
     use wasmtime_environ::prelude::*;
     use wasmtime_fiber::{Fiber, Suspend};
 
     type ContinuationFiber = Fiber<'static, &'static mut Instance, u32, ()>;
+    pub type FiberStack = wasmtime_fiber::FiberStack;
     type Yield = Suspend<&'static mut Instance, u32, ()>;
 
     /// The baseline VM continuation record.
@@ -360,7 +358,7 @@ pub mod baseline {
     /// Allocates a new continuation in suspended mode.
     #[inline(always)]
     pub fn cont_new(
-        _instance: &mut Instance,
+        instance: &mut Instance,
         func: *mut u8,
         param_count: usize,
         result_count: usize,
@@ -369,7 +367,8 @@ pub mod baseline {
         let mut values: Vec<u128> = Vec::with_capacity(capacity);
 
         let fiber = {
-            let stack = wasmtime_fiber::FiberStack::new(DEFAULT_FIBER_SIZE)
+            let stack = instance
+                .wasmfx_allocate_stack()
                 .map_err(|error| TrapReason::user_without_backtrace(error.into()))?;
             let fiber = match unsafe { func.cast::<VMFuncRef>().as_ref() } {
                 None => Fiber::new(stack, |_instance: &mut Instance, _suspend: &mut Yield| {
@@ -508,10 +507,11 @@ pub mod baseline {
 
     /// Deallocates a gives continuation reference.
     #[inline(always)]
-    pub fn drop_continuation_reference(_instance: &mut Instance, contref: *mut VMContRef) {
+    pub fn drop_continuation_reference(instance: &mut Instance, contref: *mut VMContRef) {
         // Note that continuation objects do not own their parents, so
         // we let the parent object leak.
         let contref: Box<VMContRef> = unsafe { Box::from_raw(contref) };
+        instance.wasmfx_deallocate_stack(contref.fiber.stack());
         let _: Box<ContinuationFiber> = contref.fiber;
         let _: Vec<u128> = contref.args;
         let _: Vec<u128> = contref.values;

crates/wasmtime/src/runtime/vm/fibre/unix.rs

@@ -108,15 +108,22 @@ use wasmtime_continuations::SwitchDirection;
 
 use crate::runtime::vm::{VMContext, VMFuncRef, VMOpaqueContext, ValRaw};
 
+#[derive(Debug)]
+pub enum Allocator {
+    Malloc,
+    Mmap,
+    Custom,
+}
+
 #[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,
+    // allocation strategy
+    allocator: Allocator,
 }
 
 impl FiberStack {
@@ -149,7 +156,7 @@ impl FiberStack {
             Ok(Self {
                 top: mmap.cast::<u8>().add(mmap_len),
                 len: mmap_len,
-                mmap: true,
+                allocator: Allocator::Mmap,
             })
         }
     }
@@ -158,7 +165,11 @@ impl FiberStack {
         unsafe {
             let layout = Layout::array::<u8>(size).unwrap();
             let base = alloc(layout);
-            FiberStack::from_raw_parts(base, size)
+            Ok(Self {
+                top: base.add(size),
+                len: size,
+                allocator: Allocator::Malloc,
+            })
         }
     }
 
@@ -167,7 +178,7 @@ impl FiberStack {
         Ok(Self {
             top: base.add(len),
             len,
-            mmap: false,
+            allocator: Allocator::Custom,
         })
     }
 
@@ -184,12 +195,16 @@ impl FiberStack {
 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);
-                debug_assert!(ret.is_ok());
-            } else {
-                let layout = Layout::array::<u8>(self.len).unwrap();
-                dealloc(self.top.sub(self.len), layout);
+            match self.allocator {
+                Allocator::Mmap => {
+                    let ret = rustix::mm::munmap(self.top.sub(self.len) as _, self.len);
+                    debug_assert!(ret.is_ok());
+                }
+                Allocator::Malloc => {
+                    let layout = Layout::array::<u8>(self.len).unwrap();
+                    dealloc(self.top.sub(self.len), layout);
+                }
+                Allocator::Custom => {} // It's the creator's responsibility to reclaim the memory.
             }
         }
     }

crates/wasmtime/src/runtime/vm/instance.rs

@@ -40,6 +40,8 @@ mod allocator;
 
 pub use allocator::*;
 
+use allocator::wasmfx_allocator::WasmFXAllocator;
+
 /// A type that roughly corresponds to a WebAssembly instance, but is also used
 /// for host-defined objects.
 ///
@@ -149,6 +151,9 @@ pub struct Instance {
     #[cfg(feature = "wmemcheck")]
     pub(crate) wmemcheck_state: Option<Wmemcheck>,
 
+    /// WasmFX allocator
+    wasmfx_allocator: Option<Box<WasmFXAllocator>>,
+
     /// Additional context used by compiled wasm code. This field is last, and
     /// represents a dynamically-sized array that extends beyond the nominal
     /// end of the struct (similar to a flexible array member).
@@ -189,6 +194,7 @@ impl Instance {
                 tables,
                 dropped_elements,
                 dropped_data,
+                wasmfx_allocator: None,
                 host_state: req.host_state,
                 vmctx_self_reference: SendSyncPtr::new(NonNull::new(ptr.add(1).cast()).unwrap()),
                 vmctx: VMContext {
@@ -1322,6 +1328,31 @@ impl Instance {
             self.vmctx_plus_offset_mut(self.offsets().vmctx_typed_continuations_stack_chain())
         }
     }
+
+    #[allow(dead_code)]
+    pub(crate) fn set_typed_continuations_stack_chain(&mut self, chain: *mut *mut StackChainCell) {
+        unsafe {
+            let ptr =
+                self.vmctx_plus_offset_mut(self.offsets().vmctx_typed_continuations_stack_chain());
+            *ptr = chain;
+        }
+    }
+
+    // TODO
+    pub(crate) fn wasmfx_allocate_stack(&mut self) -> Result<wasmfx_allocator::FiberStack, Error> {
+        match &self.wasmfx_allocator {
+            Some(allocator) => allocator.allocate(),
+            None => {
+                let wasmfx_config = unsafe { &*(*self.store()).wasmfx_config() };
+                self.wasmfx_allocator = Some(Box::new(WasmFXAllocator::new(wasmfx_config)?));
+                self.wasmfx_allocator.as_ref().unwrap().allocate()
+            }
+        }
+    }
+
+    pub(crate) fn wasmfx_deallocate_stack(&mut self, stack: &wasmfx_allocator::FiberStack) {
+        self.wasmfx_allocator.as_ref().unwrap().deallocate(stack)
+    }
 }
 
 /// A handle holding an `Instance` of a WebAssembly module.

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

@@ -35,6 +35,8 @@ pub use self::pooling::{
     PoolingInstanceAllocatorConfig,
 };
 
+pub mod wasmfx_allocator;
+
 /// Represents a request for a new runtime instance.
 pub struct InstanceAllocationRequest<'a> {
     /// The info related to the compiled version of this module,

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

@@ -19,7 +19,7 @@
 //! [`stack_pool`]. See those modules for more details.
 
 mod decommit_queue;
-mod index_allocator;
+pub(crate) mod index_allocator;
 mod memory_pool;
 mod table_pool;
 
@@ -81,7 +81,7 @@ use wasmtime_environ::{
     StaticModuleIndex,
 };
 
-fn round_up_to_pow2(n: usize, to: usize) -> usize {
+pub(crate) fn round_up_to_pow2(n: usize, to: usize) -> usize {
     debug_assert!(to > 0);
     debug_assert!(to.is_power_of_two());
     (n + to - 1) & !(to - 1)