diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index 9445f852a2..ae84a6203b 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -754,12 +754,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 { - debug_assert!(self.is_async()); + let config = self.inner.engine.config(); - // TODO: allocation of a fiber should be much more abstract where we - // shouldn't be allocating huge stacks on every async wasm function call. + debug_assert!(self.is_async()); + debug_assert!(config.async_stack_size > 0); + + 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?; @@ -777,18 +779,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 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 allocate 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<'_> { @@ -845,15 +875,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()); } } }