use crate::prelude::*; use crate::store::{AsStoreOpaque, Asyncness, Executor, StoreId, StoreOpaque}; use crate::vm::mpk::{self, ProtectionMask}; use crate::vm::{AlwaysMut, AsyncWasmCallState}; use crate::{Engine, StoreContextMut}; use core::mem; use core::ops::Range; use core::pin::Pin; use core::ptr::{self, NonNull}; use core::task::{Context, Poll}; use wasmtime_fiber::{Fiber, FiberStack, Suspend}; type WasmtimeResume = Result>>; type WasmtimeYield = StoreFiberYield; type WasmtimeComplete = Result<()>; type WasmtimeSuspend = Suspend; type WasmtimeFiber<'a> = Fiber<'a, WasmtimeResume, WasmtimeYield, WasmtimeComplete>; /// State related to asynchronous computations stored within a `Store`. /// /// This structure resides inside of a `Store` and is used to manage the /// various pieces of state associated with asynchronous computations. Chiefly /// this manages the `WasmtimeSuspend` pointer as well as `&mut Context<'_>` /// when polling futures. This serves as storage to use these pointers across a /// WebAssembly function boundary, for example, where the values cannot /// otherwise be explicitly threaded through. pub(crate) struct AsyncState { /// The `Suspend` for the current fiber (or null if no such fiber is /// running). /// /// This pointer is provided by the `wasmtime_fiber` crate when a fiber /// first starts, but this pointer is unable to be carried through /// WebAssembly frames for example. This serves as an alternative storage /// location for the pointer provided by `wasmtime_fiber` within a fiber's /// execution. /// /// This pointer is null when a fiber is not executing, but it is also null /// when a `BlockingContext` is created. Note that when a fiber is suspended /// it's always through a `BlockingContext` so this field is null whenever a /// fiber is suspended as well. Fiber resumption will save the prior value /// in a store and then set it to null, where suspension will then restore /// what was previously in the store. current_suspend: Option>, /// The `Context` pointer last provided in `Future for FiberFuture`. /// /// Like `current_suspend` above this is an example of a piece of context /// which needs to be carried over a WebAssembly function frame which /// otherwise doesn't take this as a parameter. This differs from /// `current_suspend` though in that it is provided as part of a `Future` /// poll operation but is "gone" after that poll operation completes. That /// means that while `current_suspend` is the same for the lifetime of a /// future this field is always changing. /// /// Like `current_suspend` though this is null either when a fiber isn't /// running or when a `BlockingContext` is created (in which case this is /// "take"en). That means that this is null on suspension/resumption of a /// fiber. /// /// The value for this pointer is threaded directly through the /// `WasmtimeResume` type which is how a pointer flows into this field from /// a future-related poll call. This means that the `BlockingContext` /// creation may take one value of a pointer here but restore another. That /// would represent suspending in one call to `Future::poll` and then /// resuming some time later in a different call to `Future::poll`. /// /// # Safety /// /// Note that this is a pretty unsafe field for two reasons. One is that /// it's a raw pointer to a `Context` provided ephemerally to some call to /// `Future::poll` on the stack. Another reason is that the lifetime /// parameter of `Context` is unsafely changed to `'static` here which is /// not correct. The ephemeral nature of this pointer is managed through the /// take-style operations in `BlockingContext` and the `'static` lifetime is /// handled by ensuring the signatures that work with `BlockingContext` all /// use constrained anonymous lifetimes that are guaranteed to be shorter /// than the original `Context` lifetime. current_future_cx: Option>>, /// The last fiber stack that was in use by the store. /// /// We use this to cache and reuse stacks as a performance optimization. // TODO: With stack switching and the Component Model Async ABI, there may // be multiple concurrent fibers in play; consider caching more than one // stack at a time and making the number tunable via `Config`. last_fiber_stack: Option, /// Whether or not this store has async host functions defined somewhere /// within it or some other store-related configuration (e.g. epochs /// yielding) which requires that wasm is executed on a fiber, thus async /// entrypoints are required. pub(crate) async_required: bool, } // SAFETY: it's known that `std::task::Context` is neither `Send` nor `Sync`, // but despite this the storage here is purely temporary in getting these // pointers across function frames. The actual types are not sent across threads // as when a store isn't polling anything the pointer values are all set to // `None`. Thus if a store is being sent across threads that's done because no // fibers are active, and once fibers are active everything will stick within // the same thread. unsafe impl Send for AsyncState {} unsafe impl Sync for AsyncState {} impl Default for AsyncState { fn default() -> Self { Self { current_suspend: None, current_future_cx: None, last_fiber_stack: None, async_required: false, } } } impl AsyncState { pub(crate) fn last_fiber_stack(&mut self) -> &mut Option { &mut self.last_fiber_stack } /// Returns whether `block_on` will succeed or panic. #[inline] pub(crate) fn can_block(&mut self) -> bool { self.current_future_cx.is_some() } } /// A helper structure used to block a fiber. /// /// This is acquired via either `StoreContextMut::with_blocking` or /// `StoreOpaque::with_blocking`. This structure represents the "taken" state of /// pointers from a store's `AsyncState`, then modeling them as safe pointers. /// /// Note that the lifetimes here are carefully controlled in instances of this /// structure through the construction of the `with` function. pub(crate) struct BlockingContext<'a, 'b> { /// Pointer to `wasmtime_fiber::Suspend` which was supplied when a fiber /// first started. /// /// When a `BlockingContext` is first created this pointer is "taken" from /// the store (the store is null'd out) and then the raw pointer previously /// in the store is unsafely transformed to this safe pointer. This /// represents how a `BlockingContext` temporarily has access to this /// suspend but when the `BlockingContext` goes away this'll make its way /// back into the store. suspend: &'a mut WasmtimeSuspend, /// Pointer to the future `Context` that this fiber is being polled with. /// /// Similar to `suspend` above this is taken from a store when a /// `BlockingContext` is created and it's restored when the /// `BlockingContext` goes away. Note though that unlike `suspend`, as /// alluded to in the documentation on `AsyncState`, this value changes over /// time as calls to poll are made. This field becomes `None` during a /// suspension because that means that the context is released and no longer /// available. Upon resumption the context here is *optionally* provided. /// Cancellation is a case where it isn't passed back and a re-poll is a /// case where it's passed back. future_cx: Option<&'a mut Context<'b>>, } impl<'a, 'b> BlockingContext<'a, 'b> { /// Method to go from a `store` provided (which internally contains a /// `StoreOpaque`) to a `BlockingContext`. /// /// This function will "take" context from `store`'s `AsyncState` field. It /// will then construct a `BlockingContext` and yield it to the closure `f` /// provided. The closure can then block on futures, suspend, etc. /// /// Upon return of the closure `f` the state from `BlockingContext` is /// restored within the store. The return value of `f` is the return value /// of this function. /// /// Note that the `store` must be provided to this function as an argument /// to originally acquire state from `AsyncState`. This store is then /// supplied back to the closure `f` provided here so the store can be used /// to construct an asynchronous or blocking computation which the /// `BlockingContext` tries to block on. /// /// # Safety /// /// This method is safe to call at any time, but it's worth noting that the /// safety of this function relies on the signature of this function. /// Notably the lifetime parameters of `BlockingContext` in the `f` closure /// here must be anonymous. That ensures that the `BlockingContext` that /// callers get access to cannot be persisted outside of that closure call /// and everything is scoped to just the closure `f` provided with nothing /// escaping. fn with(store: &mut S, f: impl FnOnce(&mut S, &mut BlockingContext<'_, '_>) -> R) -> R where S: AsStoreOpaque, { let opaque = store.as_store_opaque(); let state = opaque.fiber_async_state_mut(); // SAFETY: this is taking pointers from `AsyncState` and then unsafely // turning them into safe references. Lifetime-wise this should be safe // because the inferred lifetimes for all these pointers is constrained // by the signature of `f` provided here. That ensures that everything // is scoped purely to the closure `f` and nothing should be persisted // outside of this function call. This, for example, ensures that the // `Context<'static>` doesn't leak out, it's only with an anonymous // lifetime that's forcibly shorter. // // Provenance-wise this should be safe as if these fields in the store // are non-null then the pointers are provided up-the-stack on this // fiber and for this fiber. The "take" pattern here ensures that if // this `BlockingContext` context acquires the pointers then there are // no other instances of these pointers in use anywhere else. let future_cx = unsafe { Some(state.current_future_cx.take().unwrap().as_mut()) }; let suspend = unsafe { state.current_suspend.take().unwrap().as_mut() }; let mut reset = ResetBlockingContext { store, cx: BlockingContext { future_cx, suspend }, }; return f(&mut reset.store, &mut reset.cx); struct ResetBlockingContext<'a, 'b, S: AsStoreOpaque> { store: &'a mut S, cx: BlockingContext<'a, 'b>, } impl Drop for ResetBlockingContext<'_, '_, S> { fn drop(&mut self) { let store = self.store.as_store_opaque(); let state = store.fiber_async_state_mut(); debug_assert!(state.current_future_cx.is_none()); debug_assert!(state.current_suspend.is_none()); state.current_suspend = Some(NonNull::from(&mut *self.cx.suspend)); if let Some(cx) = &mut self.cx.future_cx { // SAFETY: while this is changing the lifetime to `'static` // it should never be used while it's `'static` given this // `BlockingContext` abstraction. state.current_future_cx = Some(NonNull::from(unsafe { change_context_lifetime(cx) })); } } } } /// Blocks on the asynchronous computation represented by `future` and /// produces the result here, in-line. /// /// This function is designed to only work when it's currently executing on /// a native fiber. This fiber provides the ability for us to handle the /// future's `Pending` state as "jump back to whomever called the fiber in /// an asynchronous fashion and propagate `Pending`". This tight coupling /// with `on_fiber` below is what powers the asynchronicity of calling wasm. /// /// This function takes a `future` and will (appear to) synchronously wait /// on the result. While this function is executing it will fiber switch /// to-and-from the original frame calling `on_fiber` which should be a /// guarantee due to how async stores are configured. /// /// The return value here is either the output of the future `T`, or a trap /// which represents that the asynchronous computation was cancelled. It is /// not recommended to catch the trap and try to keep executing wasm, so /// we've tried to liberally document this. /// /// Note that this function suspends (if needed) with /// `StoreFiberYield::KeepStore`, indicating that the store must not be used /// (and that no other fibers may be resumed) until this fiber resumes. /// Therefore, it is not appropriate for use in e.g. guest calls to /// async-lowered imports implemented as host functions, since it will /// prevent any other tasks from being run. Use `Instance::suspend` to /// suspend and release the store to allow other tasks to run before this /// fiber is resumed. /// /// # Return Value /// /// A return value of `Ok(value)` means that the future completed with /// `value`. A return value of `Err(e)` means that the fiber and its future /// have been cancelled and the fiber needs to exit and complete ASAP. /// /// # Safety /// /// This function is safe to call at any time but relies on a trait bound /// that is manually placed here the compiler does not otherwise require. /// Notably the `Send` bound on the future provided here is not required /// insofar as things compile without that. The purpose of this, however, is /// to make the `unsafe impl Send for StoreFiber` more safe. The `future` /// here is state that is stored on the stack during the suspension of this /// fiber and is otherwise not visible to the compiler. By having a `Send` /// bound here it ensures that the future doesn't have things like `Rc` or /// similar pointing into thread locals which would not be sound if this /// fiber crosses threads. pub(crate) fn block_on(&mut self, future: F) -> Result where F: Future + Send, { let mut future = core::pin::pin!(future); loop { match future.as_mut().poll(self.future_cx.as_mut().unwrap()) { Poll::Ready(v) => break Ok(v), Poll::Pending => self.suspend(StoreFiberYield::KeepStore)?, } } } /// Suspend this fiber with `yield_` as the reason. /// /// This function will suspend the current fiber and only return after the /// fiber has resumed. This function return `Ok(())` if the fiber was /// resumed to be completed, and `Err(e)` indicates that the fiber has been /// cancelled and needs to exit/complete ASAP. pub(crate) fn suspend(&mut self, yield_: StoreFiberYield) -> Result<()> { // Over a suspension point we're guaranteed that the `Context` provided // here is no longer valid, so discard it. If we're supposed to be able // to poll afterwards this will be given back as part of the resume // value given back. self.future_cx.take(); let mut new_future_cx: NonNull> = self.suspend.suspend(yield_)?; // SAFETY: this function is unsafe as we're doing "funky" things to the // `new_future_cx` we have been given. The safety here relies on the // fact that the lifetimes of `BlockingContext` are all "smaller" than // the original `Context` itself, and that should be guaranteed through // the exclusive constructor of this type `BlockingContext::with`. unsafe { self.future_cx = Some(change_context_lifetime(new_future_cx.as_mut())); } Ok(()) } } impl StoreContextMut<'_, T> { /// Blocks on the future computed by `f`. /// /// # Panics /// /// Panics if this is invoked outside the context of a fiber. #[cfg(feature = "component-model")] pub(crate) fn block_on( self, f: impl FnOnce(StoreContextMut<'_, T>) -> Pin + Send + '_>>, ) -> Result { self.with_blocking(|store, cx| cx.block_on(f(store).as_mut())) } /// Creates a `BlockingContext` suitable for blocking on futures or /// suspending the current fiber. /// /// # Panics /// /// Panics if this is invoked outside the context of a fiber. pub(crate) fn with_blocking( self, f: impl FnOnce(StoreContextMut<'_, T>, &mut BlockingContext<'_, '_>) -> R, ) -> R { BlockingContext::with(self.0, |store, cx| f(StoreContextMut(store), cx)) } } impl StoreOpaque { /// Creates a `BlockingContext` suitable for blocking on futures or /// suspending the current fiber. /// /// # Panics /// /// Panics if this is invoked outside the context of a fiber. pub(crate) fn with_blocking( &mut self, f: impl FnOnce(&mut Self, &mut BlockingContext<'_, '_>) -> R, ) -> R { BlockingContext::with(self, |store, cx| f(store, cx)) } /// Used when any configuration option that affects a store, as a side /// effect, disallows further use of sync APIs in Wasmtime. /// /// For example enabling async yielding epochs, async yielding fuel, or /// async resource limiters all require that wasm is invoked on fibers. /// These options, when enabled, will all set this flag. /// /// Note that this specifically only models the transition from "some /// previous state" to "async is now required". There's no reasonable way to /// iterate through a store and recompute this if epoch settings, for /// example, are dynamically changed. pub(crate) fn set_async_required(&mut self, asyncness: Asyncness) { match asyncness { Asyncness::Yes => { self.fiber_async_state_mut().async_required = true; } Asyncness::No => {} } } } /// Indicates whether or not a fiber needs to retain exclusive access to its /// store across a suspend/resume interval. pub(crate) enum StoreFiberYield { /// Indicates the fiber needs to retain exclusive access, meaning the store /// should not be used outside of the fiber until after the fiber either /// suspends with `ReleaseStore` or resolves. KeepStore, /// Indicates the fiber does _not_ need exclusive access across the /// suspend/resume interval, meaning the store may be used as needed until /// the fiber is resumed. #[cfg(feature = "component-model-async")] ReleaseStore, } pub(crate) struct StoreFiber<'a> { /// The raw `wasmtime_fiber::Fiber`. /// /// Note that using `StoreFiberYield` as the `Yield` type parameter allows /// the fiber to indicate whether it needs exclusive access to the store /// across suspend points (in which case it will pass `KeepStore` when /// suspending , meaning the store must not be used at all until the fiber /// is resumed again) or whether it is giving up exclusive access (in which /// case it will pass `ReleaseStore` when yielding, meaning exclusive access /// may be given to another fiber that runs concurrently. /// /// Note also that every `StoreFiber` is implicitly granted exclusive access /// to the store when it is resumed. fiber: Option>>, /// See `FiberResumeState` state: Option>, /// The Wasmtime `Engine` to which this fiber belongs. engine: Engine, /// The id of the store with which this fiber was created. /// /// Any attempt to resume a fiber with a different store than the one with /// which it was created will panic. id: StoreId, } struct RawFiber<'a>(WasmtimeFiber<'a>); impl<'a> StoreFiber<'a> { /// Convenience method to peel off some layers of abstraction around the raw /// `wasmtime_fiber::Fiber`. fn fiber(&mut self) -> Option<&mut WasmtimeFiber<'a>> { Some(&mut self.fiber.as_mut()?.get_mut().0) } /// Convenience method take the internal fiber and consume it, yielding its /// original stack. fn take_fiber_stack(&mut self) -> Option { self.fiber.take().map(|f| f.into_inner().0.into_stack()) } pub(crate) fn dispose(&mut self, store: &mut StoreOpaque) { if let Some(fiber) = self.fiber() { if !fiber.done() { let result = resume_fiber(store, self, Err(format_err!("future dropped"))); debug_assert!(result.is_ok()); } } } } // Note that this implementation will panic if the fiber is in-progress, which // will abort the process if there is already a panic being unwound. That // should only happen if we failed to call `StoreFiber::dispose` on the // in-progress fiber prior to dropping it, which indicates a bug in this crate // which must be fixed. impl Drop for StoreFiber<'_> { fn drop(&mut self) { if self.fiber.is_none() { return; } assert!( self.fiber().unwrap().done(), "attempted to drop in-progress fiber without first calling `StoreFiber::dispose`" ); self.state.take().unwrap().into_inner().dispose(); unsafe { let stack = self.take_fiber_stack().unwrap(); self.engine.allocator().deallocate_fiber_stack(stack); } } } // This is surely the most dangerous `unsafe impl Send` in the entire // crate. There are two members in `StoreFiber` which cause it to not be // `Send`. One is `suspend` and is entirely uninteresting. This is just used to // manage `Suspend` when resuming, and requires raw pointers to get it to happen // easily. Nothing too weird about the `Send`-ness, values aren't actually // crossing threads. // // The really interesting piece is `fiber`. Now the "fiber" here is actual // honest-to-god Rust code which we're moving around. What we're doing is the // equivalent of moving our thread's stack to another OS thread. Turns out we, // in general, have no idea what's on the stack and would generally have no way // to verify that this is actually safe to do! // // Thankfully, though, Wasmtime has the power. Without being glib it's actually // worth examining what's on the stack. It's unfortunately not super-local to // this function itself. Our closure to `Fiber::new` runs `func`, which is given // to us from the outside. Thankfully, though, we have tight control over // this. Usage of `on_fiber` or `Instance::resume_fiber` is typically done // *just* before entering WebAssembly itself, so we'll have a few stack frames // of Rust code (all in Wasmtime itself) before we enter wasm. // // Once we've entered wasm, well then we have a whole bunch of wasm frames on // the stack. We've got this nifty thing called Cranelift, though, which allows // us to also have complete control over everything on the stack! // // Finally, when wasm switches back to the fiber's starting pointer (this future // we're returning) then it means wasm has reentered Rust. Suspension can only // happen via either `block_on` or `Instance::suspend`. This, conveniently, also // happens entirely in Wasmtime controlled code! // // There's an extremely important point that should be called out here. // User-provided futures **are not on the stack** during suspension points. This // is extremely crucial because we in general cannot reason about Send/Sync for // stack-local variables since rustc doesn't analyze them at all. With our // construction, though, we are guaranteed that Wasmtime owns all stack frames // between the stack of a fiber and when the fiber suspends (and it could move // across threads). At this time the only user-provided piece of data on the // stack is the future itself given to us. Lo-and-behold as you might notice the // future is required to be `Send`! // // What this all boils down to is that we, as the authors of Wasmtime, need to // be extremely careful that on the async fiber stack we only store Send // things. For example we can't start using `Rc` willy nilly by accident and // leave a copy in TLS somewhere. (similarly we have to be ready for TLS to // change while we're executing wasm code between suspension points). // // While somewhat onerous it shouldn't be too too hard (the TLS bit is the // hardest bit so far). This does mean, though, that no user should ever have to // worry about the `Send`-ness of Wasmtime. If rustc says it's ok, then it's ok. // // With all that in mind we unsafely assert here that Wasmtime is correct. We // declare the fiber as only containing Send data on its stack, despite not // knowing for sure at compile time that this is correct. That's what `unsafe` // in Rust is all about, though, right? unsafe impl Send for RawFiber<'_> {} /// State of the world when a fiber last suspended. /// /// This structure represents global state that a fiber clobbers during its /// execution. For example TLS variables are updated, system resources like MPK /// masks are updated, etc. The purpose of this structure is to track all of /// this state and appropriately save/restore it around fiber suspension points. struct FiberResumeState { /// Saved list of `CallThreadState` activations that are stored on a fiber /// stack. /// /// This is a linked list that references stack-stored nodes on the fiber /// stack that is currently suspended. The `AsyncWasmCallState` type /// documents this more thoroughly but the general gist is that when we this /// fiber is resumed this linked list needs to be pushed on to the current /// thread's linked list of activations. tls: crate::runtime::vm::AsyncWasmCallState, /// Saved MPK protection mask, if enabled. /// /// When MPK is enabled then executing WebAssembly will modify the /// processor's current mask of addressable protection keys. This means that /// our current state may get clobbered when a fiber suspends. To ensure /// that this function preserves context it will, when MPK is enabled, save /// the current mask when this function is called and then restore the mask /// when the function returns (aka the fiber suspends). mpk: Option, /// The current wasm stack limit, if in use. /// /// This field stores the old of `VMStoreContext::stack_limit` that this /// fiber should be using during its execution. This is saved/restored when /// a fiber is suspended/resumed to ensure that when there are multiple /// fibers within the store they all maintain an appropriate fiber-relative /// stack limit. stack_limit: usize, /// The executor (e.g. the Pulley interpreter state) belonging to this /// fiber. /// /// This is swapped with `StoreOpaque::executor` whenever this fiber is /// resumed, suspended, or resolved. executor: Executor, } impl FiberResumeState { unsafe fn replace( self, store: &mut StoreOpaque, fiber: &mut StoreFiber<'_>, ) -> PriorFiberResumeState { let tls = unsafe { self.tls.push() }; let mpk = swap_mpk_states(self.mpk); let async_guard_range = fiber .fiber() .unwrap() .stack() .guard_range() .unwrap_or(ptr::null_mut()..ptr::null_mut()); let mut executor = self.executor; store.swap_executor(&mut executor); PriorFiberResumeState { tls, mpk, executor, stack_limit: store.replace_stack_limit(self.stack_limit), async_guard_range: store.replace_async_guard_range(async_guard_range), // The current suspend/future_cx are always null upon resumption, so // insert null. Save the old values through to get preserved across // this resume/suspend. current_suspend: store.replace_current_suspend(None), current_future_cx: store.replace_current_future_cx(None), } } fn dispose(self) { self.tls.assert_null(); } } impl StoreOpaque { /// Helper function to swap the `stack_limit` field in the `VMStoreContext` /// within this store. fn replace_stack_limit(&mut self, stack_limit: usize) -> usize { mem::replace( &mut self.vm_store_context_mut().stack_limit.get_mut(), stack_limit, ) } /// Helper function to swap the `async_guard_range` field in the `VMStoreContext` /// within this store. fn replace_async_guard_range(&mut self, range: Range<*mut u8>) -> Range<*mut u8> { mem::replace(&mut self.vm_store_context_mut().async_guard_range, range) } fn replace_current_suspend( &mut self, ptr: Option>, ) -> Option> { mem::replace(&mut self.fiber_async_state_mut().current_suspend, ptr) } fn replace_current_future_cx( &mut self, ptr: Option>>, ) -> Option>> { mem::replace(&mut self.fiber_async_state_mut().current_future_cx, ptr) } } struct PriorFiberResumeState { tls: crate::runtime::vm::PreviousAsyncWasmCallState, mpk: Option, stack_limit: usize, async_guard_range: Range<*mut u8>, current_suspend: Option>, current_future_cx: Option>>, executor: Executor, } impl PriorFiberResumeState { unsafe fn replace(self, store: &mut StoreOpaque) -> FiberResumeState { let tls = unsafe { self.tls.restore() }; let mpk = swap_mpk_states(self.mpk); // No need to save `_my_guard` since we can re-infer it from the fiber // that this state is attached to. let _my_guard = store.replace_async_guard_range(self.async_guard_range); // Restore the previous values of current_{suspend,future_cx} but we // should be guaranteed that the prior values are null, so double-check // that here. let prev = store.replace_current_suspend(self.current_suspend); assert!(prev.is_none()); let prev = store.replace_current_future_cx(self.current_future_cx); assert!(prev.is_none()); let mut executor = self.executor; store.swap_executor(&mut executor); FiberResumeState { tls, mpk, executor, stack_limit: store.replace_stack_limit(self.stack_limit), } } } fn swap_mpk_states(mask: Option) -> Option { mask.map(|mask| { let current = mpk::current_mask(); mpk::allow(mask); current }) } /// Resume the specified fiber, granting it exclusive access to the store with /// which it was created. /// /// This will return `Ok(result)` if the fiber resolved, where `result` is the /// returned value; it will return `Err(yield_)` if the fiber suspended, where /// `yield_` indicates whether it released access to the store or not. See /// `StoreFiber::fiber` for details. fn resume_fiber<'a>( store: &mut StoreOpaque, fiber: &mut StoreFiber<'a>, result: WasmtimeResume, ) -> Result { assert_eq!(store.id(), fiber.id); struct Restore<'a, 'b> { store: &'b mut StoreOpaque, fiber: &'b mut StoreFiber<'a>, state: Option, } impl Drop for Restore<'_, '_> { fn drop(&mut self) { self.fiber.state = Some(unsafe { self.state.take().unwrap().replace(self.store).into() }); } } let result = unsafe { let prev = fiber .state .take() .unwrap() .into_inner() .replace(store, fiber); let restore = Restore { store, fiber, state: Some(prev), }; restore.fiber.fiber().unwrap().resume(result) }; match &result { // The fiber has finished, so recycle its stack by disposing of the // underlying fiber itself. Ok(_) => { if let Some(stack) = fiber.take_fiber_stack() { store.deallocate_fiber_stack(stack); } } // The fiber has not yet finished, so it stays as-is. Err(_) => { // If `Err` is returned that means the fiber suspended, so we // propagate that here. // // An additional safety check is performed when leaving this // function to help bolster the guarantees of `unsafe impl Send` // above. Notably this future may get re-polled on a different // thread. Wasmtime's thread-local state points to the stack, // however, meaning that it would be incorrect to leave a pointer in // TLS when this function returns. This function performs a runtime // assert to verify that this is the case, notably that the one TLS // pointer Wasmtime uses is not pointing anywhere within the // stack. If it is then that's a bug indicating that TLS management // in Wasmtime is incorrect. if let Some(range) = fiber.fiber().unwrap().stack().range() { AsyncWasmCallState::assert_current_state_not_in_range(range); } } } result } /// Create a new `StoreFiber` which runs the specified closure. /// /// # Safety /// /// The returned `StoreFiber<'a>` structure is unconditionally `Send` but the /// send-ness is actually a function of `S`. When `S` is statically known to be /// `Send` then use the safe [`make_fiber`] function. pub(crate) unsafe fn make_fiber_unchecked<'a, S>( store: &mut S, fun: impl FnOnce(&mut S) -> Result<()> + Send + Sync + 'a, ) -> Result> where S: AsStoreOpaque + ?Sized + 'a, { let opaque = store.as_store_opaque(); let engine = opaque.engine().clone(); let executor = Executor::new(&engine)?; let id = opaque.id(); let stack = opaque.allocate_fiber_stack()?; let track_pkey_context_switch = opaque.has_pkey(); let store = &raw mut *store; let fiber = Fiber::new(stack, move |result: WasmtimeResume, suspend| { let future_cx = match result { Ok(cx) => cx, // Cancelled before we started? Just return. Err(_) => return Ok(()), }; // SAFETY: This fiber will only be resumed using `resume_fiber`, which // takes a `&mut StoreOpaque` parameter and has given us exclusive // access to the store until we exit or yield it back to the resumer. let store_ref = unsafe { &mut *store }; // It should be a guarantee that the store has null pointers here upon // starting a fiber, so now's the time to fill in the pointers now that // the fiber is running and `future_cx` and `suspend` are both in scope. // Note that these pointers are removed when this function returns as // that's when they fall out of scope. let async_state = store_ref.as_store_opaque().fiber_async_state_mut(); assert!(async_state.current_suspend.is_none()); assert!(async_state.current_future_cx.is_none()); async_state.current_suspend = Some(NonNull::from(suspend)); async_state.current_future_cx = Some(future_cx); struct ResetCurrentPointersToNull<'a, S>(&'a mut S) where S: AsStoreOpaque + ?Sized; impl Drop for ResetCurrentPointersToNull<'_, S> where S: AsStoreOpaque + ?Sized, { fn drop(&mut self) { let state = self.0.as_store_opaque().fiber_async_state_mut(); // Double-check that the current suspension isn't null (it // should be what's in this closure). Note though that we // can't check `current_future_cx` because it may either be // here or not be here depending on whether this was // cancelled or not. debug_assert!(state.current_suspend.is_some()); state.current_suspend = None; state.current_future_cx = None; } } let reset = ResetCurrentPointersToNull(store_ref); fun(reset.0) })?; Ok(StoreFiber { state: Some( FiberResumeState { tls: crate::runtime::vm::AsyncWasmCallState::new(), mpk: if track_pkey_context_switch { Some(ProtectionMask::all()) } else { None }, stack_limit: usize::MAX, executor, } .into(), ), engine, id, fiber: Some(RawFiber(fiber).into()), }) } /// Safe wrapper around [`make_fiber_unchecked`] which requires that `S` is /// `Send`. #[cfg(feature = "component-model-async")] pub(crate) fn make_fiber<'a, S>( store: &mut S, fun: impl FnOnce(&mut S) -> Result<()> + Send + Sync + 'a, ) -> Result> where S: AsStoreOpaque + Send + ?Sized + 'a, { unsafe { make_fiber_unchecked(store, fun) } } /// Run the specified function on a newly-created fiber and `.await` its /// completion. pub(crate) async fn on_fiber( store: &mut S, func: impl FnOnce(&mut S) -> R + Send + Sync, ) -> Result where S: AsStoreOpaque + ?Sized, R: Send + Sync, { let opaque = store.as_store_opaque(); let config = opaque.engine().config(); debug_assert!(config.async_stack_size > 0); let mut result = None; // SAFETY: the `StoreFiber` returned by `make_fiber_unchecked` is `Send` // despite we not actually knowing here whether `S` is `Send` or not. That // is safe here, however, because this function is already conditionally // `Send` based on `S`. Additionally `fiber` doesn't escape this function, // so the future-of-this-function is still correctly `Send`-vs-not. let fiber = unsafe { make_fiber_unchecked(store, |store| { result = Some(func(store)); Ok(()) })? }; { let fiber = FiberFuture { store: store.as_store_opaque(), fiber: Some(fiber), #[cfg(feature = "component-model-async")] on_release: OnRelease::ReturnPending, } .await .unwrap(); debug_assert!(fiber.is_none()); } Ok(result.unwrap()) } /// Run the specified fiber until it either suspends with /// `StoreFiberYield::ReleaseStore` or resolves. /// /// This will return `Some` if the fiber suspends with /// `StoreFiberYield::ReleaseStore` or else `None` if it resolves. #[cfg(feature = "component-model-async")] pub(crate) async fn resolve_or_release<'a>( store: &mut StoreOpaque, fiber: StoreFiber<'a>, ) -> Result>> { FiberFuture { store, fiber: Some(fiber), on_release: OnRelease::ReturnReady, } .await } /// Tells a `FiberFuture` what to do if `poll_fiber` returns /// `Err(StoreFiberYield::ReleaseStore)`. #[cfg(feature = "component-model-async")] enum OnRelease { /// Return `Poll::Pending` from `FiberFuture::poll` ReturnPending, /// Return `Poll::Ready` from `FiberFuture::poll`, handing ownership of the /// `StoreFiber` to the caller. ReturnReady, } /// A `Future` implementation for running a `StoreFiber` to completion, giving /// it exclusive access to its store until it resolves. struct FiberFuture<'a, 'b> { store: &'a mut StoreOpaque, fiber: Option>, #[cfg(feature = "component-model-async")] on_release: OnRelease, } impl<'b> Future for FiberFuture<'_, 'b> { type Output = Result>>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let me = self.get_mut(); // SAFETY: We need to carry over this `cx` into our fiber's runtime for // when it tries to poll sub-futures that are created. Doing this must // be done unsafely, however, since `cx` is only alive for this one // singular function call. Here we do a `transmute` to extend the // lifetime of `Context` so it can be stored in our `Store`, and then we // replace the current polling context with this one. // // The safety of this extension relies on never actually using // `Context<'static>` with `'static` actually there, which should be // satisfied by the users of this in the `BlockingContext` structure // where the lifetime parameters there are always more constrained than // they are here. let cx: &mut Context<'static> = unsafe { change_context_lifetime(cx) }; let cx = NonNull::from(cx); match resume_fiber(me.store, me.fiber.as_mut().unwrap(), Ok(cx)) { Ok(Ok(())) => Poll::Ready(Ok(None)), Ok(Err(e)) => Poll::Ready(Err(e)), Err(StoreFiberYield::KeepStore) => Poll::Pending, #[cfg(feature = "component-model-async")] Err(StoreFiberYield::ReleaseStore) => match &me.on_release { OnRelease::ReturnPending => Poll::Pending, OnRelease::ReturnReady => Poll::Ready(Ok(me.fiber.take())), }, } } } impl Drop for FiberFuture<'_, '_> { fn drop(&mut self) { if let Some(fiber) = &mut self.fiber { fiber.dispose(self.store); } } } /// Changes the lifetime `'l` in `Context<'l>` to something else. /// /// # Safety /// /// Not a safe operation. Requires external knowledge about how the pointer is /// being used to determine whether it's actually safe or not. See docs on /// callers of this function. The purpose of this is to scope the `transmute` to /// as small an operation as possible. unsafe fn change_context_lifetime<'a, 'b>(cx: &'a mut Context<'_>) -> &'a mut Context<'b> { // SAFETY: See the function documentation, this is not safe in general. unsafe { mem::transmute::<&mut Context<'_>, &mut Context<'b>>(cx) } }