diff --git a/crates/fiber/src/lib.rs b/crates/fiber/src/lib.rs index d17f92d6ba29..0dee07a2b933 100644 --- a/crates/fiber/src/lib.rs +++ b/crates/fiber/src/lib.rs @@ -21,6 +21,14 @@ cfg_if::cfg_if! { /// Represents an execution stack to use for a fiber. pub struct FiberStack(imp::FiberStack); +fn _assert_send_sync() { + fn _assert_send() {} + fn _assert_sync() {} + + _assert_send::(); + _assert_sync::(); +} + impl FiberStack { /// Creates a new fiber stack of the given size. pub fn new(size: usize) -> io::Result { diff --git a/crates/fiber/src/unix.rs b/crates/fiber/src/unix.rs index 11a7a55e1839..a9aa3104442c 100644 --- a/crates/fiber/src/unix.rs +++ b/crates/fiber/src/unix.rs @@ -36,7 +36,7 @@ use std::ops::Range; use std::ptr; pub struct FiberStack { - base: *mut u8, + base: BasePtr, len: usize, /// Stored here to ensure that when this `FiberStack` the backing storage, @@ -44,6 +44,11 @@ pub struct FiberStack { storage: FiberStackStorage, } +struct BasePtr(*mut u8); + +unsafe impl Send for BasePtr {} +unsafe impl Sync for BasePtr {} + enum FiberStackStorage { Mmap(#[allow(dead_code)] MmapFiberStack), Unmanaged(usize), @@ -64,7 +69,7 @@ impl FiberStack { // region so the base and length of our stack are both offset by a // single page. Ok(FiberStack { - base: stack.mapping_base.wrapping_byte_add(page_size), + base: BasePtr(stack.mapping_base.wrapping_byte_add(page_size)), len: stack.mapping_len - page_size, storage: FiberStackStorage::Mmap(stack), }) @@ -77,7 +82,7 @@ impl FiberStack { return Self::from_custom(asan::new_fiber_stack(len)?); } Ok(FiberStack { - base: base.add(guard_size), + base: BasePtr(base.add(guard_size)), len, storage: FiberStackStorage::Unmanaged(guard_size), }) @@ -101,28 +106,28 @@ impl FiberStack { "expected fiber stack end ({end_ptr:?}) to be page aligned ({page_size:#x})", ); Ok(FiberStack { - base: start_ptr, + base: BasePtr(start_ptr), len: range.len(), storage: FiberStackStorage::Custom(custom), }) } pub fn top(&self) -> Option<*mut u8> { - Some(self.base.wrapping_byte_add(self.len)) + Some(self.base.0.wrapping_byte_add(self.len)) } pub fn range(&self) -> Option> { - let base = self.base as usize; + let base = self.base.0 as usize; Some(base..base + self.len) } pub fn guard_range(&self) -> Option> { match &self.storage { FiberStackStorage::Unmanaged(guard_size) => unsafe { - let start = self.base.sub(*guard_size); - Some(start..self.base) + let start = self.base.0.sub(*guard_size); + Some(start..self.base.0) }, - FiberStackStorage::Mmap(mmap) => Some(mmap.mapping_base..self.base), + FiberStackStorage::Mmap(mmap) => Some(mmap.mapping_base..self.base.0), FiberStackStorage::Custom(custom) => Some(custom.guard_range()), } } diff --git a/crates/wasmtime/src/runtime/store.rs b/crates/wasmtime/src/runtime/store.rs index ac22f08407ba..7948e7ce2264 100644 --- a/crates/wasmtime/src/runtime/store.rs +++ b/crates/wasmtime/src/runtime/store.rs @@ -335,6 +335,7 @@ pub struct StoreOpaque { table_limit: usize, #[cfg(feature = "async")] async_state: AsyncState, + // If fuel_yield_interval is enabled, then we store the remaining fuel (that isn't in // runtime_limits) here. The total amount of fuel is the runtime limits and reserve added // together. Then when we run out of gas, we inject the yield amount from the reserve @@ -392,6 +393,8 @@ pub struct StoreOpaque { struct AsyncState { current_suspend: UnsafeCell<*mut wasmtime_fiber::Suspend, (), Result<()>>>, current_poll_cx: UnsafeCell, + /// The last fiber stack that was in use by this store. + last_fiber_stack: Option, } #[cfg(feature = "async")] @@ -556,6 +559,7 @@ impl Store { async_state: AsyncState { current_suspend: UnsafeCell::new(ptr::null_mut()), current_poll_cx: UnsafeCell::new(PollContext::default()), + last_fiber_stack: None, }, fuel_reserve: 0, fuel_yield_interval: None, @@ -2099,6 +2103,31 @@ at https://bytecodealliance.org/security. core::ptr::null_mut()..core::ptr::null_mut() } } + + #[cfg(feature = "async")] + fn allocate_fiber_stack(&mut self) -> Result { + if let Some(stack) = self.async_state.last_fiber_stack.take() { + return Ok(stack); + } + self.engine().allocator().allocate_fiber_stack() + } + + #[cfg(feature = "async")] + fn deallocate_fiber_stack(&mut self, stack: wasmtime_fiber::FiberStack) { + self.flush_fiber_stack(); + self.async_state.last_fiber_stack = Some(stack); + } + + /// Releases the last fiber stack to the underlying instance allocator, if + /// present. + fn flush_fiber_stack(&mut self) { + #[cfg(feature = "async")] + if let Some(stack) = self.async_state.last_fiber_stack.take() { + unsafe { + self.engine.allocator().deallocate_fiber_stack(stack); + } + } + } } impl StoreContextMut<'_, T> { @@ -2124,13 +2153,14 @@ impl StoreContextMut<'_, T> { debug_assert!(config.async_stack_size > 0); let mut slot = None; - let future = { + let mut future = { let current_poll_cx = self.0.async_state.current_poll_cx.get(); let current_suspend = self.0.async_state.current_suspend.get(); - let stack = self.engine().allocator().allocate_fiber_stack()?; + let stack = self.0.allocate_fiber_stack()?; let engine = self.engine().clone(); let slot = &mut slot; + let this = &mut *self; let fiber = wasmtime_fiber::Fiber::new(stack, move |keep_going, suspend| { // First check and see if we were interrupted/dropped, and only // continue if we haven't been. @@ -2148,7 +2178,7 @@ impl StoreContextMut<'_, T> { let _reset = Reset(current_suspend, *current_suspend); *current_suspend = suspend; - *slot = Some(func(self)); + *slot = Some(func(this)); Ok(()) } })?; @@ -2163,7 +2193,12 @@ impl StoreContextMut<'_, T> { state: Some(crate::runtime::vm::AsyncWasmCallState::new()), } }; - future.await?; + (&mut future).await?; + let stack = future.fiber.take().map(|f| f.into_stack()); + drop(future); + if let Some(stack) = stack { + self.0.deallocate_fiber_stack(stack); + } return Ok(slot.unwrap()); @@ -2373,6 +2408,10 @@ impl StoreContextMut<'_, T> { // completion. impl Drop for FiberFuture<'_> { fn drop(&mut self) { + if self.fiber.is_none() { + return; + } + if !self.fiber().done() { let result = self.resume(Err(anyhow!("future dropped"))); // This resumption with an error should always complete the @@ -2737,6 +2776,8 @@ impl fmt::Debug for Store { impl Drop for Store { fn drop(&mut self) { + self.inner.flush_fiber_stack(); + // for documentation on this `unsafe`, see `into_data`. unsafe { ManuallyDrop::drop(&mut self.inner.data);