allow recursive Wasm invocation from concurrent host functions (#12152)

* allow recursive Wasm invocation from concurrent host functions

The core changes here are:

- remove an unnecessary assertion from `concurrent::prepare_call`
- track instance states (e.g. backpressure, etc.) on a per `(Instance, RuntimeComponentInstanceIndex)` basis
    - both parts of that key are needed now that concurrent state is tracked on a per-store basis rather than a per-instance basis since `RuntimeComponentInstanceIndex`es are not globally unique

While discussing the above with Alex, we realized the use of a `HashMap` to
track per-instance states was both pessimal and unnecessary.  Consequently, I've
folded that state into `ComponentInstance::instance_handle_tables`, renaming it
to `instance_states`.  That involved a fair amount of code churn, but doesn't
change behavior except as described in the second bullet point above.

Thanks to Alex for the test case!

Fixes #12098

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

* use new `RuntimeInstance` type instead of tuples

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

---------

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

Author: dicej

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

@@ -2492,7 +2492,7 @@ impl<T> StoreContextMut<'_, T> {
                     *producer.lock().unwrap() = Some((mine, buffer));
 
                     if write_buffer {
-                        write( token, id, producer.clone(), kind).await?;
+                        write(token, id, producer.clone(), kind).await?;
                     }
 
                     Ok(if dropped {
@@ -4261,21 +4261,21 @@ impl Instance {
 
 impl ComponentInstance {
     fn table_for_transmit(self: Pin<&mut Self>, ty: TransmitIndex) -> &mut HandleTable {
-        let (tables, types) = self.guest_tables();
+        let (states, types) = self.instance_states();
         let runtime_instance = match ty {
             TransmitIndex::Stream(ty) => types[ty].instance,
             TransmitIndex::Future(ty) => types[ty].instance,
         };
-        &mut tables[runtime_instance]
+        states[runtime_instance].handle_table()
     }
 
     fn table_for_error_context(
         self: Pin<&mut Self>,
         ty: TypeComponentLocalErrorContextTableIndex,
     ) -> &mut HandleTable {
-        let (tables, types) = self.guest_tables();
+        let (states, types) = self.instance_states();
         let runtime_instance = types[ty].instance;
-        &mut tables[runtime_instance]
+        states[runtime_instance].handle_table()
     }
 
     fn get_mut_by_index(
@@ -4456,7 +4456,8 @@ impl Waitable {
             } => {
                 let instance = instance.id().get_mut(store);
                 let runtime_instance = instance.component().types()[ty].instance;
-                let (rep, state) = instance.guest_tables().0[runtime_instance]
+                let (rep, state) = instance.instance_states().0[runtime_instance]
+                    .handle_table()
                     .future_rep(ty, handle)
                     .unwrap();
                 assert_eq!(rep, self.rep());
@@ -4477,7 +4478,8 @@ impl Waitable {
             } => {
                 let instance = instance.id().get_mut(store);
                 let runtime_instance = instance.component().types()[ty].instance;
-                let (rep, state) = instance.guest_tables().0[runtime_instance]
+                let (rep, state) = instance.instance_states().0[runtime_instance]
+                    .handle_table()
                     .stream_rep(ty, handle)
                     .unwrap();
                 assert_eq!(rep, self.rep());

crates/wasmtime/src/runtime/component/concurrent/futures_and_streams.rs

@@ -794,7 +794,7 @@ impl Func {
             ResourceTables {
                 host_table: Some(host_table),
                 calls,
-                guest: Some(instance.guest_tables()),
+                guest: Some(instance.instance_states()),
             }
             .exit_call()?;
         }

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

@@ -286,7 +286,7 @@ impl<'a, T: 'static> LowerContext<'a, T> {
             ResourceTables {
                 host_table: Some(host_table),
                 calls,
-                guest: Some(instance.guest_tables()),
+                guest: Some(instance.instance_states()),
             },
             host_resource_data,
         )
@@ -473,7 +473,7 @@ impl<'a> LiftContext<'a> {
                 calls: self.calls,
                 // Note that the unsafety here should be valid given the contract of
                 // `LiftContext::new`.
-                guest: Some(self.instance.as_mut().guest_tables()),
+                guest: Some(self.instance.as_mut().instance_states()),
             },
             self.host_resource_data,
         )

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

@@ -659,7 +659,7 @@ fn resource_tables<'a>(
     ResourceTables {
         host_table: None,
         calls,
-        guest: Some(instance.guest_tables()),
+        guest: Some(instance.instance_states()),
     }
 }
 

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

@@ -49,7 +49,7 @@ impl ComponentStoreData {
     }
 
     #[cfg(feature = "component-model-async")]
-    pub(crate) fn assert_guest_tables_empty(&mut self) {
+    pub(crate) fn assert_instance_states_empty(&mut self) {
         for (_, instance) in self.instances.iter_mut() {
             let Some(instance) = instance.as_mut() else {
                 continue;
@@ -58,10 +58,11 @@ impl ComponentStoreData {
             assert!(
                 instance
                     .get_mut()
-                    .guest_tables()
+                    .instance_states()
                     .0
-                    .iter()
-                    .all(|(_, table)| table.is_empty())
+                    .iter_mut()
+                    .all(|(_, state)| state.handle_table().is_empty()
+                        && state.concurrent_state().pending_is_empty())
             );
         }
     }

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

@@ -9,6 +9,8 @@
 use crate::component::{Component, Instance, InstancePre, ResourceType, RuntimeImport};
 use crate::module::ModuleRegistry;
 use crate::runtime::component::ComponentInstanceId;
+#[cfg(feature = "component-model-async")]
+use crate::runtime::component::concurrent::ConcurrentInstanceState;
 use crate::runtime::vm::instance::{InstanceLayout, OwnedInstance, OwnedVMContext};
 use crate::runtime::vm::vmcontext::VMFunctionBody;
 use crate::runtime::vm::{
@@ -45,6 +47,35 @@ pub use self::handle_table::{TransmitLocalState, Waitable};
 pub use self::resources::CallContext;
 pub use self::resources::{CallContexts, ResourceTables, TypedResource, TypedResourceIndex};
 
+/// Represents the state of a (sub-)component instance.
+#[derive(Default)]
+pub struct InstanceState {
+    /// Represents the Component Model Async state of a (sub-)component instance.
+    #[cfg(feature = "component-model-async")]
+    concurrent_state: ConcurrentInstanceState,
+
+    /// State of handles (e.g. resources, waitables, etc.) for this instance.
+    ///
+    /// For resource handles, this is paired with other information to create a
+    /// `ResourceTables` and manipulated through that.  For other handles, this
+    /// is used directly to translate guest handles to host representations and
+    /// vice-versa.
+    handle_table: HandleTable,
+}
+
+impl InstanceState {
+    /// Represents the Component Model Async state of a (sub-)component instance.
+    #[cfg(feature = "component-model-async")]
+    pub fn concurrent_state(&mut self) -> &mut ConcurrentInstanceState {
+        &mut self.concurrent_state
+    }
+
+    /// State of handles (e.g. resources, waitables, etc.) for this instance.
+    pub fn handle_table(&mut self) -> &mut HandleTable {
+        &mut self.handle_table
+    }
+}
+
 /// Runtime representation of a component instance and all state necessary for
 /// the instance itself.
 ///
@@ -86,13 +117,9 @@ pub struct ComponentInstance {
     // borrowing a store mutably at the same time as a contained instance.
     component: Component,
 
-    /// State of handles (e.g. resources, waitables, etc.) for this component.
-    ///
-    /// For resource handles, this is paired with other information to create a
-    /// `ResourceTables` and manipulated through that.  For other handles, this
-    /// is used directly to translate guest handles to host representations and
-    /// vice-versa.
-    instance_handle_tables: PrimaryMap<RuntimeComponentInstanceIndex, HandleTable>,
+    /// Contains state specific to each (sub-)component instance within this
+    /// top-level instance.
+    instance_states: PrimaryMap<RuntimeComponentInstanceIndex, InstanceState>,
 
     /// What all compile-time-identified core instances are mapped to within the
     /// `Store` that this component belongs to.
@@ -293,16 +320,15 @@ impl ComponentInstance {
     ) -> OwnedComponentInstance {
         let offsets = VMComponentOffsets::new(HostPtr, component.env_component());
         let num_instances = component.env_component().num_runtime_component_instances;
-        let mut instance_handle_tables =
-            PrimaryMap::with_capacity(num_instances.try_into().unwrap());
+        let mut instance_states = PrimaryMap::with_capacity(num_instances.try_into().unwrap());
         for _ in 0..num_instances {
-            instance_handle_tables.push(HandleTable::default());
+            instance_states.push(InstanceState::default());
         }
 
         let mut ret = OwnedInstance::new(ComponentInstance {
             id,
             offsets,
-            instance_handle_tables,
+            instance_states,
             instances: PrimaryMap::with_capacity(
                 component
                     .env_component()
@@ -825,22 +851,30 @@ impl ComponentInstance {
         resource_instance == component.defined_resource_instances[idx]
     }
 
-    /// Returns the runtime state of resources associated with this component.
+    /// Returns the runtime state of resources and concurrency associated with
+    /// this component.
     #[inline]
-    pub fn guest_tables(
+    pub fn instance_states(
         self: Pin<&mut Self>,
     ) -> (
-        &mut PrimaryMap<RuntimeComponentInstanceIndex, HandleTable>,
+        &mut PrimaryMap<RuntimeComponentInstanceIndex, InstanceState>,
         &ComponentTypes,
     ) {
         // safety: we've chosen the `pin` guarantee of `self` to not apply to
         // the map returned.
         unsafe {
             let me = self.get_unchecked_mut();
-            (&mut me.instance_handle_tables, me.component.types())
+            (&mut me.instance_states, me.component.types())
         }
     }
 
+    pub fn instance_state(
+        self: Pin<&mut Self>,
+        instance: RuntimeComponentInstanceIndex,
+    ) -> Option<&mut InstanceState> {
+        self.instance_states().0.get_mut(instance)
+    }
+
     /// Returns the destructor and instance flags for the specified resource
     /// table type.
     ///

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

@@ -5,7 +5,7 @@ use crate::component::Instance;
 use crate::component::concurrent::WaitResult;
 use crate::prelude::*;
 #[cfg(feature = "component-model-async")]
-use crate::runtime::component::concurrent::ResourcePair;
+use crate::runtime::component::concurrent::{ResourcePair, RuntimeInstance};
 use crate::runtime::vm::component::{ComponentInstance, VMComponentContext};
 use crate::runtime::vm::{HostResultHasUnwindSentinel, VMStore, VmSafe};
 use core::cell::Cell;
@@ -672,12 +672,15 @@ fn trap(_store: &mut dyn VMStore, _instance: Instance, code: u8) -> Result<()> {
 #[cfg(feature = "component-model-async")]
 fn backpressure_modify(
     store: &mut dyn VMStore,
-    _instance: Instance,
+    instance: Instance,
     caller_instance: u32,
     increment: u8,
 ) -> Result<()> {
-    store.concurrent_state_mut().backpressure_modify(
-        RuntimeComponentInstanceIndex::from_u32(caller_instance),
+    store.backpressure_modify(
+        RuntimeInstance {
+            instance: instance.id().instance(),
+            index: RuntimeComponentInstanceIndex::from_u32(caller_instance),
+        },
         |old| {
             if increment != 0 {
                 old.checked_add(1)

crates/wasmtime/src/runtime/vm/component/libcalls.rs

@@ -20,7 +20,7 @@
 //! about ABI details can be found in lifting/lowering throughout Wasmtime,
 //! namely in the `Resource<T>` and `ResourceAny` types.
 
-use super::{HandleTable, RemovedResource};
+use super::{HandleTable, InstanceState, RemovedResource};
 use crate::prelude::*;
 use core::error::Error;
 use core::fmt;
@@ -53,7 +53,7 @@ pub struct ResourceTables<'a> {
     /// `ResourceAny::resource_drop` which won't consult this table as it's
     /// only operating over the host table.
     pub guest: Option<(
-        &'a mut PrimaryMap<RuntimeComponentInstanceIndex, HandleTable>,
+        &'a mut PrimaryMap<RuntimeComponentInstanceIndex, InstanceState>,
         &'a ComponentTypes,
     )>,
 
@@ -196,8 +196,8 @@ impl ResourceTables<'_> {
         match resource {
             TypedResource::Host(_) => self.host_table.as_mut().unwrap(),
             TypedResource::Component { ty, .. } => {
-                let (tables, types) = self.guest.as_mut().unwrap();
-                &mut tables[types[*ty].unwrap_concrete_instance()]
+                let (states, types) = self.guest.as_mut().unwrap();
+                states[types[*ty].unwrap_concrete_instance()].handle_table()
             }
         }
     }
@@ -206,8 +206,8 @@ impl ResourceTables<'_> {
         match index {
             TypedResourceIndex::Host(_) => self.host_table.as_mut().unwrap(),
             TypedResourceIndex::Component { ty, .. } => {
-                let (tables, types) = self.guest.as_mut().unwrap();
-                &mut tables[types[*ty].unwrap_concrete_instance()]
+                let (states, types) = self.guest.as_mut().unwrap();
+                states[types[*ty].unwrap_concrete_instance()].handle_table()
             }
         }
     }

crates/wasmtime/src/runtime/vm/component/resources.rs

@@ -767,3 +767,73 @@ async fn cancel_host_future() -> Result<()> {
         }
     }
 }
+
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn run_wasm_in_call_async() -> Result<()> {
+    _ = env_logger::try_init();
+
+    let mut config = Config::new();
+    config.async_support(true);
+    config.wasm_component_model_async(true);
+    let engine = Engine::new(&config)?;
+
+    let a = Component::new(
+        &engine,
+        r#"
+(component
+  (type $t (func async))
+  (import "a" (func $f (type $t)))
+  (core func $f (canon lower (func $f)))
+  (core module $a
+    (import "" "f" (func $f))
+    (func (export "run") call $f)
+  )
+  (core instance $a (instantiate $a
+    (with "" (instance (export "f" (func $f))))
+  ))
+  (func (export "run") (type $t)
+    (canon lift (core func $a "run")))
+)
+        "#,
+    )?;
+    let b = Component::new(
+        &engine,
+        r#"
+(component
+  (type $t (func async))
+  (core module $a
+    (func (export "run"))
+  )
+  (core instance $a (instantiate $a))
+  (func (export "run") (type $t)
+    (canon lift (core func $a "run")))
+)
+        "#,
+    )?;
+
+    type State = Option<Instance>;
+
+    let mut linker = Linker::new(&engine);
+    linker
+        .root()
+        .func_wrap_concurrent("a", |accessor: &Accessor<State>, (): ()| {
+            Box::pin(async move {
+                let func = accessor.with(|mut access| {
+                    access
+                        .get()
+                        .unwrap()
+                        .get_typed_func::<(), ()>(&mut access, "run")
+                })?;
+                func.call_concurrent(accessor, ()).await?;
+                Ok(())
+            })
+        })?;
+    let mut store = Store::new(&engine, None);
+    let instance_a = linker.instantiate_async(&mut store, &a).await?;
+    let instance_b = linker.instantiate_async(&mut store, &b).await?;
+    *store.data_mut() = Some(instance_b);
+    let run = instance_a.get_typed_func::<(), ()>(&mut store, "run")?;
+    run.call_async(&mut store, ()).await?;
+    Ok(())
+}