From 697f4c012aba35e82e2d97fad1b645fac84f169d Mon Sep 17 00:00:00 2001 From: Peter Huene Date: Thu, 4 Feb 2021 13:05:01 -0800 Subject: [PATCH] Implement allocating fiber stacks for an instance allocator. This commit implements allocating fiber stacks in an instance allocator. The on-demand instance allocator doesn't support custom stacks, so the implementation will use the allocation from `wasmtime-fiber` for the fiber stacks. In the future, the pooling instance allocator will return custom stacks to use on Linux and macOS. On Windows, the native fiber implementation will always be used. --- crates/fiber/src/lib.rs | 21 ++++++++ crates/fiber/src/unix.rs | 61 ++++++++++++++-------- crates/fiber/src/windows.rs | 12 ++++- crates/runtime/Cargo.toml | 6 +++ crates/runtime/src/instance/allocator.rs | 41 +++++++++++++++ crates/runtime/src/lib.rs | 3 ++ crates/wasmtime/Cargo.toml | 2 +- crates/wasmtime/src/config.rs | 59 +++++++++++++++++---- crates/wasmtime/src/store.rs | 66 +++++++++++++++++++----- 9 files changed, 223 insertions(+), 48 deletions(-) diff --git a/crates/fiber/src/lib.rs b/crates/fiber/src/lib.rs index bafae2f01c13..6c835c511f81 100644 --- a/crates/fiber/src/lib.rs +++ b/crates/fiber/src/lib.rs @@ -51,6 +51,27 @@ impl<'a, Resume, Yield, Return> Fiber<'a, Resume, Yield, Return> { }) } + /// Creates a new fiber with existing stack space that will execute `func`. + /// + /// This function returns a `Fiber` which, when resumed, will execute `func` + /// to completion. When desired the `func` can suspend itself via + /// `Fiber::suspend`. + /// + /// # Safety + /// + /// The caller must properly allocate the stack space with a guard page and + /// make the pages accessible for correct behavior. + pub unsafe fn new_with_stack( + top_of_stack: *mut u8, + func: impl FnOnce(Resume, &Suspend) -> Return + 'a, + ) -> io::Result> { + Ok(Fiber { + inner: imp::Fiber::new_with_stack(top_of_stack, func), + done: Cell::new(false), + _phantom: PhantomData, + }) + } + /// Resumes execution of this fiber. /// /// This function will transfer execution to the fiber and resume from where diff --git a/crates/fiber/src/unix.rs b/crates/fiber/src/unix.rs index 2c1069041ebe..c14a188f50f5 100644 --- a/crates/fiber/src/unix.rs +++ b/crates/fiber/src/unix.rs @@ -35,10 +35,10 @@ use std::io; use std::ptr; pub struct Fiber { - // Description of the mmap region we own. This should be abstracted - // eventually so we aren't personally mmap-ing this region. - mmap: *mut libc::c_void, - mmap_len: usize, + // The top of the stack; for stacks allocated by the fiber implementation itself, + // the base address of the allocation will be `top_of_stack.sub(alloc_len.unwrap())` + top_of_stack: *mut u8, + alloc_len: Option, } pub struct Suspend { @@ -66,21 +66,40 @@ where } impl Fiber { - pub fn new(stack_size: usize, func: F) -> io::Result + pub fn new(stack_size: usize, func: F) -> io::Result + where + F: FnOnce(A, &super::Suspend) -> C, + { + let fiber = Self::alloc_with_stack(stack_size)?; + fiber.init(func); + Ok(fiber) + } + + pub fn new_with_stack(top_of_stack: *mut u8, func: F) -> Self + where + F: FnOnce(A, &super::Suspend) -> C, + { + let fiber = Self { + top_of_stack, + alloc_len: None, + }; + + fiber.init(func); + + fiber + } + + fn init(&self, func: F) where F: FnOnce(A, &super::Suspend) -> C, { - let fiber = Fiber::alloc_with_stack(stack_size)?; unsafe { - // Initialize the top of the stack to be resumed from - let top_of_stack = fiber.top_of_stack(); let data = Box::into_raw(Box::new(func)).cast(); - wasmtime_fiber_init(top_of_stack, fiber_start::, data); - Ok(fiber) + wasmtime_fiber_init(self.top_of_stack, fiber_start::, data); } } - fn alloc_with_stack(stack_size: usize) -> io::Result { + fn alloc_with_stack(stack_size: usize) -> io::Result { unsafe { // Round up our stack size request to the nearest multiple of the // page size. @@ -104,7 +123,10 @@ impl Fiber { if mmap == libc::MAP_FAILED { return Err(io::Error::last_os_error()); } - let ret = Fiber { mmap, mmap_len }; + let ret = Self { + top_of_stack: mmap.cast::().add(mmap_len), + alloc_len: Some(mmap_len), + }; let res = libc::mprotect( mmap.cast::().add(page_size).cast(), stack_size, @@ -124,27 +146,24 @@ impl Fiber { // stack, otherwise known as our reserved slot for this information. // // In the diagram above this is updating address 0xAff8 - let top_of_stack = self.top_of_stack(); - let addr = top_of_stack.cast::().offset(-1); + let addr = self.top_of_stack.cast::().offset(-1); addr.write(result as *const _ as usize); - wasmtime_fiber_switch(top_of_stack); + wasmtime_fiber_switch(self.top_of_stack); // null this out to help catch use-after-free addr.write(0); } } - - unsafe fn top_of_stack(&self) -> *mut u8 { - self.mmap.cast::().add(self.mmap_len) - } } impl Drop for Fiber { fn drop(&mut self) { unsafe { - let ret = libc::munmap(self.mmap, self.mmap_len); - debug_assert!(ret == 0); + if let Some(alloc_len) = self.alloc_len { + let ret = libc::munmap(self.top_of_stack.sub(alloc_len) as _, alloc_len); + debug_assert!(ret == 0); + } } } } diff --git a/crates/fiber/src/windows.rs b/crates/fiber/src/windows.rs index 69a5f161e6f0..c40fb8aeb031 100644 --- a/crates/fiber/src/windows.rs +++ b/crates/fiber/src/windows.rs @@ -40,7 +40,7 @@ where } impl Fiber { - pub fn new(stack_size: usize, func: F) -> io::Result + pub fn new(stack_size: usize, func: F) -> io::Result where F: FnOnce(A, &super::Suspend) -> C, { @@ -61,11 +61,19 @@ impl Fiber { drop(Box::from_raw(state.initial_closure.get().cast::())); Err(io::Error::last_os_error()) } else { - Ok(Fiber { fiber, state }) + Ok(Self { fiber, state }) } } } + pub fn new_with_stack(_top_of_stack: *mut u8, _func: F) -> Self + where + F: FnOnce(A, &super::Suspend) -> C, + { + // Windows fibers have no support for custom stacks + unimplemented!() + } + pub(crate) fn resume(&self, result: &Cell>) { unsafe { let is_fiber = IsThreadAFiber() != 0; diff --git a/crates/runtime/Cargo.toml b/crates/runtime/Cargo.toml index 40266d379371..82c0b5fddd9a 100644 --- a/crates/runtime/Cargo.toml +++ b/crates/runtime/Cargo.toml @@ -33,3 +33,9 @@ cc = "1.0" [badges] maintenance = { status = "actively-developed" } + +[features] +default = ["async"] + +# Enables support for "async" fiber stacks in the instance allocator +async = [] diff --git a/crates/runtime/src/instance/allocator.rs b/crates/runtime/src/instance/allocator.rs index 6c16604b0764..17aa76389c5a 100644 --- a/crates/runtime/src/instance/allocator.rs +++ b/crates/runtime/src/instance/allocator.rs @@ -73,6 +73,17 @@ pub enum InstantiationError { #[error("Trap occurred during instantiation")] Trap(Trap), } +/// An error while creating a fiber stack. +#[cfg(feature = "async")] +#[derive(Error, Debug)] +pub enum FiberStackError { + /// An error for when the allocator doesn't support custom fiber stacks. + #[error("Custom fiber stacks are not supported by the allocator")] + NotSupported, + /// A limit on how many fibers are supported has been reached. + #[error("Limit of {0} concurrent fibers has been reached")] + Limit(u32), +} /// Represents a runtime instance allocator. /// @@ -127,6 +138,24 @@ pub unsafe trait InstanceAllocator: Send + Sync { /// /// Use extreme care when deallocating an instance so that there are no dangling instance pointers. unsafe fn deallocate(&self, handle: &InstanceHandle); + + /// Allocates a fiber stack for calling async functions on. + /// + /// Returns the top of the fiber stack if successfully allocated. + #[cfg(feature = "async")] + fn allocate_fiber_stack(&self) -> Result<*mut u8, FiberStackError>; + + /// Deallocates a fiber stack that was previously allocated. + /// + /// # Safety + /// + /// This function is unsafe because there are no guarantees that the given stack + /// is no longer in use. + /// + /// Additionally, passing a stack pointer that was not returned from `allocate_fiber_stack` + /// will lead to undefined behavior. + #[cfg(feature = "async")] + unsafe fn deallocate_fiber_stack(&self, stack: *mut u8); } unsafe fn initialize_vmcontext( @@ -544,4 +573,16 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator { ptr::drop_in_place(instance as *const Instance as *mut Instance); alloc::dealloc(instance as *const Instance as *mut _, layout); } + + #[cfg(feature = "async")] + fn allocate_fiber_stack(&self) -> Result<*mut u8, FiberStackError> { + // The on-demand allocator does not support allocating fiber stacks + Err(FiberStackError::NotSupported) + } + + #[cfg(feature = "async")] + unsafe fn deallocate_fiber_stack(&self, _stack: *mut u8) { + // This should never be called as `allocate_fiber_stack` never returns success + unreachable!() + } } diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs index 04ef0026ce2c..85b3aa01963b 100644 --- a/crates/runtime/src/lib.rs +++ b/crates/runtime/src/lib.rs @@ -55,6 +55,9 @@ pub use crate::vmcontext::{ VMSharedSignatureIndex, VMTableDefinition, VMTableImport, VMTrampoline, }; +#[cfg(feature = "async")] +pub use crate::instance::FiberStackError; + /// Version number of this crate. pub const VERSION: &str = env!("CARGO_PKG_VERSION"); diff --git a/crates/wasmtime/Cargo.toml b/crates/wasmtime/Cargo.toml index d063f0de04ae..92a408f1562a 100644 --- a/crates/wasmtime/Cargo.toml +++ b/crates/wasmtime/Cargo.toml @@ -71,4 +71,4 @@ experimental_x64 = ["wasmtime-jit/experimental_x64"] # Enables support for "async stores" as well as defining host functions as # `async fn` and calling functions asynchronously. -async = ["wasmtime-fiber"] +async = ["wasmtime-fiber", "wasmtime-runtime/async"] diff --git a/crates/wasmtime/src/config.rs b/crates/wasmtime/src/config.rs index 324b03501d13..19151600b54f 100644 --- a/crates/wasmtime/src/config.rs +++ b/crates/wasmtime/src/config.rs @@ -58,6 +58,8 @@ pub struct Config { pub(crate) max_instances: usize, pub(crate) max_tables: usize, pub(crate) max_memories: usize, + #[cfg(feature = "async")] + pub(crate) async_stack_size: usize, } impl Config { @@ -108,6 +110,8 @@ impl Config { max_instances: 10_000, max_tables: 10_000, max_memories: 10_000, + #[cfg(feature = "async")] + async_stack_size: 2 << 20, }; ret.wasm_backtrace_details(WasmBacktraceDetails::Environment); return ret; @@ -182,23 +186,58 @@ impl Config { self } - /// Configures the maximum amount of native stack space available to + /// Configures the maximum amount of stack space available for /// executing WebAssembly code. /// - /// WebAssembly code currently executes on the native call stack for its own - /// call frames. WebAssembly, however, also has well-defined semantics on - /// stack overflow. This is intended to be a knob which can help configure - /// how much native stack space a wasm module is allowed to consume. Note - /// that the number here is not super-precise, but rather wasm will take at - /// most "pretty close to this much" stack space. + /// WebAssembly has well-defined semantics on stack overflow. This is + /// intended to be a knob which can help configure how much stack space + /// wasm execution is allowed to consume. Note that the number here is not + /// super-precise, but rather wasm will take at most "pretty close to this + /// much" stack space. /// /// If a wasm call (or series of nested wasm calls) take more stack space /// than the `size` specified then a stack overflow trap will be raised. /// - /// By default this option is 1 MB. - pub fn max_wasm_stack(&mut self, size: usize) -> &mut Self { + /// When the `async` feature is enabled, this value cannot exceed the + /// `async_stack_size` option. Be careful not to set this value too close + /// to `async_stack_size` as doing so may limit how much stack space + /// is available for host functions. Unlike wasm functions that trap + /// on stack overflow, a host function that overflows the stack will + /// abort the process. + /// + /// By default this option is 1 MiB. + pub fn max_wasm_stack(&mut self, size: usize) -> Result<&mut Self> { + #[cfg(feature = "async")] + if size > self.async_stack_size { + bail!("wasm stack size cannot exceed the async stack size"); + } + + if size == 0 { + bail!("wasm stack size cannot be zero"); + } + self.max_wasm_stack = size; - self + Ok(self) + } + + /// Configures the size of the stacks used for asynchronous execution. + /// + /// This setting configures the size of the stacks that are allocated for + /// asynchronous execution. The value cannot be less than `max_wasm_stack`. + /// + /// The amount of stack space guaranteed for host functions is + /// `async_stack_size - max_wasm_stack`, so take care not to set these two values + /// close to one another; doing so may cause host functions to overflow the + /// stack and abort the process. + /// + /// By default this option is 2 MiB. + #[cfg(feature = "async")] + pub fn async_stack_size(&mut self, size: usize) -> Result<&mut Self> { + if size < self.max_wasm_stack { + bail!("async stack size cannot be less than the maximum wasm stack size"); + } + self.async_stack_size = size; + Ok(self) } /// Configures whether the WebAssembly threads proposal will be enabled for diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index 96afecbc648f..8c0c2763e8af 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -749,12 +749,14 @@ impl Store { /// that the various comments are illuminating as to what's going on here. #[cfg(feature = "async")] pub(crate) async fn on_fiber(&self, func: impl FnOnce() -> R) -> Result { + let config = self.inner.engine.config(); + debug_assert!(self.is_async()); + debug_assert!(config.async_stack_size > 0); - // TODO: allocation of a fiber should be much more abstract where we - // shouldn't be allocating huge stacks on every async wasm function call. + type SuspendType = wasmtime_fiber::Suspend, (), Result<(), Trap>>; let mut slot = None; - let fiber = wasmtime_fiber::Fiber::new(10 * 1024 * 1024, |keep_going, suspend| { + let func = |keep_going, suspend: &SuspendType| { // First check and see if we were interrupted/dropped, and only // continue if we haven't been. keep_going?; @@ -772,18 +774,46 @@ impl Store { slot = Some(func()); Ok(()) - }) - .map_err(|e| Trap::from(anyhow::Error::from(e)))?; + }; + + let (fiber, stack) = match config.instance_allocator().allocate_fiber_stack() { + Ok(stack) => { + // Use the returned stack and take care to deallocate it when finished + ( + unsafe { + wasmtime_fiber::Fiber::new_with_stack(stack, func) + .map_err(|e| Trap::from(anyhow::Error::from(e)))? + }, + stack, + ) + } + Err(wasmtime_runtime::FiberStackError::NotSupported) => { + // The allocator doesn't support custom fiber stacks for the current platform + // Request that the fiber itself allocates the stack + ( + wasmtime_fiber::Fiber::new(config.async_stack_size, func) + .map_err(|e| Trap::from(anyhow::Error::from(e)))?, + std::ptr::null_mut(), + ) + } + Err(e) => return Err(Trap::from(anyhow::Error::from(e))), + }; // Once we have the fiber representing our synchronous computation, we // wrap that in a custom future implementation which does the // translation from the future protocol to our fiber API. - FiberFuture { fiber, store: self }.await?; + FiberFuture { + fiber, + store: self, + stack, + } + .await?; return Ok(slot.unwrap()); struct FiberFuture<'a> { fiber: wasmtime_fiber::Fiber<'a, Result<(), Trap>, (), Result<(), Trap>>, store: &'a Store, + stack: *mut u8, } impl Future for FiberFuture<'_> { @@ -840,15 +870,23 @@ impl Store { // completion. impl Drop for FiberFuture<'_> { fn drop(&mut self) { - if self.fiber.done() { - return; + if !self.fiber.done() { + let result = self.fiber.resume(Err(Trap::new("future dropped"))); + // This resumption with an error should always complete the + // fiber. While it's technically possible for host code to catch + // the trap and re-resume, we'd ideally like to signal that to + // callers that they shouldn't be doing that. + debug_assert!(result.is_ok()); + } + if !self.stack.is_null() { + unsafe { + self.store + .engine() + .config() + .instance_allocator() + .deallocate_fiber_stack(self.stack) + }; } - let result = self.fiber.resume(Err(Trap::new("future dropped"))); - // This resumption with an error should always complete the - // fiber. While it's technically possible for host code to catch - // the trap and re-resume, we'd ideally like to signal that to - // callers that they shouldn't be doing that. - debug_assert!(result.is_ok()); } } }