generalize async fiber abstraction (#11114)

* generalize async fiber abstraction

As part of the work implementing the new Component Model async ABI in the
`wasip3-prototyping` repo, I've generalized the `FiberFuture` abstraction in
`wasmtime::runtime::store::async_` to support fibers which can either retain
exclusive access to the store across suspend points or release it.  The latter
allows the store to be used by the `component-model-async` event loop and/or
other fibers to run before the original fiber resumes, which is the key to
allowing multiple fibers to run concurrently, passing control of the store back
and forth.

In the case of Pulley, the above generalization means we also need to give each
fiber its own `Interpreter` so that multiple concurrent fibers don't clobber
each other's state.

Concretely, this moves a lot of the code out of `async_.rs` and into a new
`fiber.rs` submodule which will be shared with the `component-model-async`
implementation.

This also pulls in a new `StoreToken<T>` utility which has been useful in
`wasip3-prototyping` to safely convert from a `&mut dyn VMStore` to a
`StoreContextMut<'a, T>` when we previously witnessed a conversion in the other
direction.

Note that I've added a `'static` bound to the `VMStore` trait, which simplifies
use of `&mut dyn VMStore`, avoiding thorny lifetime issues.

Signed-off-by: Joel Dice <[email protected]>

* address review feedback

Signed-off-by: Joel Dice <[email protected]>

* fix miri-flagged stacked borrow violation

As part of my earlier effort to unify the fiber abstractions in the `wasmtime`
crate, I changed a `*mut StoreFiber` field to a `&mut StoreFiber`, not realizing
that it resulted in a mutable alias at runtime and thus undefined behavior.
Miri caught it, fortunately.

Signed-off-by: Joel Dice <[email protected]>

* remove unneeded `Send` bounds

Signed-off-by: Joel Dice <[email protected]>

* address more review feedback

Main changes:

- Make `resume_fiber[_raw]` take a `&mut StoreOpaque` parameter to make its unsafe internals easier to reason about, safety-wise.

- Panic if `StoreFiber::drop` is called on an in-progress fiber without having called `StoreFiber::dispose` to gracefully end it first.

- (Re)introduce `FiberFuture`, which closes over a `&mut StoreOpaque` and uses it to call `StoreFiber::dispose` on drop.

This will require a few more changes to make it usable by `concurrent.rs`, but
I'll save those changes for a later PR.

Signed-off-by: Joel Dice <[email protected]>

* address more review feedback

Signed-off-by: Joel Dice <[email protected]>

* update `impl Send For StoreFiber` comment

Signed-off-by: Joel Dice <[email protected]>

* Remove currently-extraneous `Result<()>` from fibers

May be needed for concurrent bits, but for now not necessary.

* Use safe pointers instead of raw pointers

It's predicted Miri won't like this, but for now in-repo it's ok with it.

* Fold more responsibility into `resume_fiber_raw`

Remove the need for the function entirely and replace it with
`resume_fiber`.

* Remove channels from async fibers

Can use stack-based closures/results to transmit the result instead of
needing a channel.

* Fold `on_fiber_raw` directly into `on_fiber`

The `on_fiber` function is small enough it should be possible to do so.

* Don't use `Option` in `FiberFuture`

Leave the fiber non-optional at-rest so it's always available for the
destructor.

* Fold `suspend` functions together

Small shims, not otherwise public at this time, so remove a layer of
indirection.

* Move stack limit management to `FiberResumeState`

Helps remove some raw pointers that are held for a long time within
`AsyncCx`

* add some doc comments to `fiber.rs`

Signed-off-by: Joel Dice <[email protected]>

* update `fiber.rs` and friends to match CM async requirements

This adds a `resolve_or_release` function, which `Instance::resume_fiber` will
use when current `concurrent.rs` stub is replaced by a real implementation.

Signed-off-by: Joel Dice <[email protected]>

* fix non-component-model-async build warnings

Signed-off-by: Joel Dice <[email protected]>

* make `resume_fiber` private in `fiber.rs`

Signed-off-by: Joel Dice <[email protected]>

* Shrink `PollContext` state

Move management of the async guard range elsewhere to the normal
save/restore area.

* Refactor `AsyncCx`, reduce `unsafe`

* Remove the `AsyncCx` type from Wasmtime as it's inherently `unsafe` to
  use, instead bundle operations directly on a `Store*` reference.

* Don't retain pointers-to-pointers within the roughly-equivalent
  `BlockingContext` created in this PR. Instead when a blocking context
  is created "take" the metadata from the store to assert exclusive
  ownership of the pointers.

* Refactor how `&mut Context<'_>` is passed around, namely thread it
  through fiber parameters to model resumption as registering a new
  context to poll with.

* Remove `PollContext` in favor of directly storing a pointer as it's
  now mostly an empty structure.

* Minor refactorings to make things more future-refactorable and/or
  clear in a few places.

* Refactor management of the "current suspend" and "current future
  context" pointers. These are now null'd out on resumption and asserted
  null on suspension.

* Remove the need for a generic `Reset` structure in the fiber bits as
  it's a pretty dangerous structure to have in general.

The end result of this refactoring is that all usage of `block_on` is
now safe and additionally many of the internals of the implementation
are safer than they were before

* Adjust some lint attributes

* Make manipulation of `AsyncState` safe

No need for raw pointers with recent refactorings.

* Fix dead code warning

* More dead code warnings

* Cut down on raw pointers in fiber.rs

* Move executor save/restore to normal fiber state save/restore

* Bikeshed a method name

* update comment in make_fiber

Signed-off-by: Joel Dice <[email protected]>

* fix machports build

Signed-off-by: Joel Dice <[email protected]>

---------

Signed-off-by: Joel Dice <[email protected]>
Co-authored-by: Alex Crichton <[email protected]>

Author: dicej

benches/call.rs

@@ -134,8 +134,8 @@ fn bench_host_to_wasm<Params, Results>(
     typed_params: Params,
     typed_results: Results,
 ) where
-    Params: WasmParams + ToVals + Copy,
-    Results: WasmResults + ToVals + Copy + PartialEq + Debug,
+    Params: WasmParams + ToVals + Copy + Sync,
+    Results: WasmResults + ToVals + Copy + Sync + PartialEq + Debug + 'static,
 {
     // Benchmark the "typed" version, which should be faster than the versions
     // below.
@@ -628,7 +628,8 @@ mod component {
             + PartialEq
             + Debug
             + Send
-            + Sync,
+            + Sync
+            + 'static,
     {
         // Benchmark the "typed" version.
         c.bench_function(&format!("component - host-to-wasm - typed - {name}"), |b| {

crates/wasmtime/src/runtime.rs

@@ -34,6 +34,8 @@ pub(crate) mod code_memory;
 #[cfg(feature = "debug-builtins")]
 pub(crate) mod debug;
 pub(crate) mod externals;
+#[cfg(feature = "async")]
+pub(crate) mod fiber;
 pub(crate) mod gc;
 pub(crate) mod instance;
 pub(crate) mod instantiate;

crates/wasmtime/src/runtime/component/func.rs

@@ -547,7 +547,7 @@ impl Func {
         let mut store = store.as_context_mut();
         assert!(
             store.0.async_support(),
-            "cannot use `call_async` without enabling async support in the config"
+            "cannot use `post_return_async` without enabling async support in the config"
         );
         // Future optimization opportunity: conditionally use a fiber here since
         // some func's post_return will not need the async context (i.e. end up

crates/wasmtime/src/runtime/component/func/typed.rs

@@ -179,7 +179,7 @@ where
     ) -> Result<Return>
     where
         Params: Send + Sync,
-        Return: Send + Sync,
+        Return: Send + Sync + 'static,
     {
         let mut store = store.as_context_mut();
         assert!(

crates/wasmtime/src/runtime/component/linker.rs

@@ -441,10 +441,8 @@ impl<T: 'static> LinkerInstance<'_, T> {
             self.engine.config().async_support,
             "cannot use `func_wrap_async` without enabling async support in the config"
         );
-        let ff = move |mut store: StoreContextMut<'_, T>, params: Params| -> Result<Return> {
-            let async_cx = store.as_context_mut().0.async_cx().expect("async cx");
-            let future = f(store.as_context_mut(), params);
-            unsafe { async_cx.block_on(Pin::from(future)) }?
+        let ff = move |store: StoreContextMut<'_, T>, params: Params| -> Result<Return> {
+            store.block_on(|store| f(store, params).into())?
         };
         self.func_wrap(name, ff)
     }
@@ -603,12 +601,10 @@ impl<T: 'static> LinkerInstance<'_, T> {
             self.engine.config().async_support,
             "cannot use `func_new_async` without enabling async support in the config"
         );
-        let ff = move |mut store: StoreContextMut<'_, T>, params: &[Val], results: &mut [Val]| {
-            let async_cx = store.as_context_mut().0.async_cx().expect("async cx");
-            let future = f(store.as_context_mut(), params, results);
-            unsafe { async_cx.block_on(Pin::from(future)) }?
+        let ff = move |store: StoreContextMut<'_, T>, params: &[Val], results: &mut [Val]| {
+            store.with_blocking(|store, cx| cx.block_on(Pin::from(f(store, params, results)))?)
         };
-        self.func_new(name, ff)
+        return self.func_new(name, ff);
     }
 
     /// Defines a [`Module`] within this instance.
@@ -676,12 +672,8 @@ impl<T: 'static> LinkerInstance<'_, T> {
         let dtor = Arc::new(crate::func::HostFunc::wrap_inner(
             &self.engine,
             move |mut cx: crate::Caller<'_, T>, (param,): (u32,)| {
-                let async_cx = cx.as_context_mut().0.async_cx().expect("async cx");
-                let future = dtor(cx.as_context_mut(), param);
-                match unsafe { async_cx.block_on(Pin::from(future)) } {
-                    Ok(Ok(())) => Ok(()),
-                    Ok(Err(trap)) | Err(trap) => Err(trap),
-                }
+                cx.as_context_mut()
+                    .block_on(|store| dtor(store, param).into())?
             },
         ));
         self.insert(name, Definition::Resource(ty, dtor))?;

crates/wasmtime/src/runtime/component/store.rs

@@ -28,6 +28,13 @@ impl ComponentStoreData {
     pub fn next_component_instance_id(&self) -> ComponentInstanceId {
         self.instances.next_key()
     }
+
+    #[cfg(feature = "component-model-async")]
+    pub(crate) fn drop_fibers(store: &mut StoreOpaque) {
+        _ = store;
+        // This function will actually do something when runtime support for
+        // `component-model-async` is merged.
+    }
 }
 
 impl StoreData {