Consolidate component-related store data

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

@@ -1805,7 +1805,7 @@ impl StoreOpaque {
         if !task.returned_or_cancelled() {
             log::trace!("push call context for {guest_task:?}");
             let call_context = task.call_context.take().unwrap();
-            self.component_resource_state().0.push(call_context);
+            self.component_call_contexts_mut().push(call_context);
         }
         Ok(())
     }
@@ -1820,7 +1820,7 @@ impl StoreOpaque {
             .returned_or_cancelled()
         {
             log::trace!("pop call context for {guest_task:?}");
-            let call_context = Some(self.component_resource_state().0.pop().unwrap());
+            let call_context = Some(self.component_call_contexts_mut().pop().unwrap());
             self.concurrent_state_mut()
                 .get_mut(guest_task)?
                 .call_context = call_context;
@@ -2745,8 +2745,7 @@ impl Instance {
                         token
                             .as_context_mut(store)
                             .0
-                            .component_resource_state()
-                            .0
+                            .component_call_contexts_mut()
                             .push(call_context);
                     }
                 });
@@ -2761,8 +2760,7 @@ impl Instance {
                             token
                                 .as_context_mut(store)
                                 .0
-                                .component_resource_state()
-                                .0
+                                .component_call_contexts_mut()
                                 .pop()
                                 .unwrap(),
                         );

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

@@ -728,7 +728,8 @@ impl<'a> Instantiator<'a> {
         imports: &'a Arc<PrimaryMap<RuntimeImportIndex, RuntimeImport>>,
     ) -> Result<Instantiator<'a>> {
         let env_component = component.env_component();
-        store.register_component(component)?;
+        let (modules, engine) = store.modules_and_engine_mut();
+        modules.register_component(component, engine)?;
         let imported_resources: ImportedResources =
             PrimaryMap::with_capacity(env_component.imported_resources.len());
 

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

@@ -1,19 +1,42 @@
+use crate::runtime::component::concurrent::ConcurrentState;
+use crate::runtime::component::{HostResourceData, Instance};
+use crate::runtime::vm;
 #[cfg(feature = "component-model-async")]
 use crate::runtime::vm::VMStore;
+use crate::runtime::vm::component::{CallContexts, HandleTable};
 use crate::runtime::vm::component::{ComponentInstance, OwnedComponentInstance};
 use crate::store::{StoreData, StoreId, StoreOpaque};
+use crate::{Engine, StoreContextMut};
 #[cfg(feature = "component-model-async")]
 use alloc::vec::Vec;
 use core::pin::Pin;
 use wasmtime_environ::PrimaryMap;
 use wasmtime_environ::component::RuntimeComponentInstanceIndex;
 
-#[derive(Default)]
+/// Extensions to `Store` which are only relevant for component-related
+/// information.
 pub struct ComponentStoreData {
+    /// All component instances, in a similar manner to how core wasm instances
+    /// are managed.
     instances: PrimaryMap<ComponentInstanceId, Option<OwnedComponentInstance>>,
 
     /// Whether an instance belonging to this store has trapped.
     trapped: bool,
+
+    /// Total number of component instances in this store, used to track
+    /// resources in the instance allocator.
+    num_component_instances: usize,
+
+    /// Runtime state for components used in the handling of resources, borrow,
+    /// and calls. These also interact with the `ResourceAny` type and its
+    /// internal representation.
+    component_host_table: HandleTable,
+    component_calls: CallContexts,
+    host_resource_data: HostResourceData,
+
+    /// Metadata/tasks/etc related to component-model-async and concurrency
+    /// support.
+    concurrent_state: Option<ConcurrentState>,
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
@@ -26,29 +49,51 @@ pub struct RuntimeInstance {
     pub index: RuntimeComponentInstanceIndex,
 }
 
-impl StoreOpaque {
-    pub(crate) fn trapped(&self) -> bool {
-        self.store_data().components.trapped
-    }
-
-    pub(crate) fn set_trapped(&mut self) {
-        self.store_data_mut().components.trapped = true;
+impl ComponentStoreData {
+    pub fn new(engine: &Engine) -> ComponentStoreData {
+        ComponentStoreData {
+            instances: Default::default(),
+            trapped: false,
+            num_component_instances: 0,
+            component_host_table: Default::default(),
+            component_calls: Default::default(),
+            host_resource_data: Default::default(),
+            concurrent_state: if engine.tunables().concurrency_support {
+                #[cfg(feature = "component-model-async")]
+                {
+                    Some(Default::default())
+                }
+                #[cfg(not(feature = "component-model-async"))]
+                {
+                    unreachable!()
+                }
+            } else {
+                None
+            },
+        }
     }
-}
 
-impl StoreData {
-    pub(crate) fn push_component_instance(
-        &mut self,
-        data: OwnedComponentInstance,
-    ) -> ComponentInstanceId {
-        let expected = data.get().id();
-        let ret = self.components.instances.push(Some(data));
-        assert_eq!(expected, ret);
-        ret
+    /// Hook used just before a `Store` is dropped to dispose of anything
+    /// necessary.
+    ///
+    /// Used at this time to deallocate fibers related to concurrency support.
+    pub fn run_manual_drop_routines<T>(store: StoreContextMut<T>) {
+        // We need to drop the fibers of each component instance before
+        // attempting to drop the instances themselves since the fibers may need
+        // to be resumed and allowed to exit cleanly before we yank the state
+        // out from under them.
+        //
+        // This will also drop any futures which might use a `&Accessor` fields
+        // in their `Drop::drop` implementations, in which case they'll need to
+        // be called from with in the context of a `tls::set` closure.
+        #[cfg(feature = "component-model-async")]
+        if store.0.component_data().concurrent_state.is_some() {
+            ComponentStoreData::drop_fibers_and_futures(store.0);
+        }
+        #[cfg(not(feature = "component-model-async"))]
+        let _ = store;
     }
-}
 
-impl ComponentStoreData {
     pub fn next_component_instance_id(&self) -> ComponentInstanceId {
         self.instances.next_key()
     }
@@ -86,32 +131,11 @@ impl ComponentStoreData {
             );
         }
     }
-}
-
-impl StoreData {
-    pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
-        self.components.instances[id].as_ref().unwrap().get()
-    }
 
-    pub(crate) fn component_instance_mut(
-        &mut self,
-        id: ComponentInstanceId,
-    ) -> Pin<&mut ComponentInstance> {
-        self.components.instances[id].as_mut().unwrap().get_mut()
-    }
-}
-
-impl StoreOpaque {
-    pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
-        self.store_data().component_instance(id)
-    }
-
-    #[cfg(feature = "component-model-async")]
-    pub(crate) fn component_instance_mut(
-        &mut self,
-        id: ComponentInstanceId,
-    ) -> Pin<&mut ComponentInstance> {
-        self.store_data_mut().component_instance_mut(id)
+    pub fn decrement_allocator_resources(&mut self, allocator: &dyn vm::InstanceAllocator) {
+        for _ in 0..self.num_component_instances {
+            allocator.decrement_component_instance_count();
+        }
     }
 }
 
@@ -194,11 +218,149 @@ impl StoreComponentInstanceId {
     }
 
     /// Same as `get_mut`, but borrows less of a store.
-    pub(crate) fn from_data_get_mut<'a>(
-        &self,
-        store: &'a mut StoreData,
-    ) -> Pin<&'a mut ComponentInstance> {
+    fn from_data_get_mut<'a>(&self, store: &'a mut StoreData) -> Pin<&'a mut ComponentInstance> {
         self.assert_belongs_to(store.id());
         store.component_instance_mut(self.instance)
     }
 }
+
+impl StoreData {
+    pub(crate) fn push_component_instance(
+        &mut self,
+        data: OwnedComponentInstance,
+    ) -> ComponentInstanceId {
+        let expected = data.get().id();
+        let ret = self.components.instances.push(Some(data));
+        assert_eq!(expected, ret);
+        ret
+    }
+
+    pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
+        self.components.instances[id].as_ref().unwrap().get()
+    }
+
+    pub(crate) fn component_instance_mut(
+        &mut self,
+        id: ComponentInstanceId,
+    ) -> Pin<&mut ComponentInstance> {
+        self.components.instances[id].as_mut().unwrap().get_mut()
+    }
+}
+
+impl StoreOpaque {
+    pub(crate) fn trapped(&self) -> bool {
+        self.store_data().components.trapped
+    }
+
+    pub(crate) fn set_trapped(&mut self) {
+        self.store_data_mut().components.trapped = true;
+    }
+
+    pub(crate) fn component_data(&self) -> &ComponentStoreData {
+        &self.store_data().components
+    }
+
+    pub(crate) fn component_data_mut(&mut self) -> &mut ComponentStoreData {
+        &mut self.store_data_mut().components
+    }
+
+    pub(crate) fn component_call_contexts_mut(&mut self) -> &mut CallContexts {
+        &mut self.component_data_mut().component_calls
+    }
+
+    pub(crate) fn push_component_instance(&mut self, instance: Instance) {
+        // We don't actually need the instance itself right now, but it seems
+        // like something we will almost certainly eventually want to keep
+        // around, so force callers to provide it.
+        let _ = instance;
+
+        self.component_data_mut().num_component_instances += 1;
+    }
+
+    pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
+        self.store_data().component_instance(id)
+    }
+
+    #[cfg(feature = "component-model-async")]
+    pub(crate) fn component_instance_mut(
+        &mut self,
+        id: ComponentInstanceId,
+    ) -> Pin<&mut ComponentInstance> {
+        self.store_data_mut().component_instance_mut(id)
+    }
+
+    #[cfg(feature = "component-model-async")]
+    pub(crate) fn concurrent_state_mut(&mut self) -> &mut ConcurrentState {
+        debug_assert!(self.concurrency_support());
+        self.component_data_mut().concurrent_state.as_mut().unwrap()
+    }
+
+    #[inline]
+    pub(crate) fn concurrency_support(&self) -> bool {
+        let support = self.component_data().concurrent_state.is_some();
+        debug_assert_eq!(support, self.engine().tunables().concurrency_support);
+        support
+    }
+
+    pub(crate) fn component_resource_state_with_instance_and_concurrent_state(
+        &mut self,
+        instance: Instance,
+    ) -> (
+        &mut CallContexts,
+        &mut HandleTable,
+        &mut HostResourceData,
+        Pin<&mut ComponentInstance>,
+        Option<&mut ConcurrentState>,
+    ) {
+        let instance = instance.id();
+        instance.assert_belongs_to(self.id());
+        let data = self.component_data_mut();
+        (
+            &mut data.component_calls,
+            &mut data.component_host_table,
+            &mut data.host_resource_data,
+            data.instances[instance.instance]
+                .as_mut()
+                .unwrap()
+                .get_mut(),
+            data.concurrent_state.as_mut(),
+        )
+    }
+
+    pub(crate) fn component_resource_tables(
+        &mut self,
+        instance: Option<Instance>,
+    ) -> vm::component::ResourceTables<'_> {
+        self.component_resource_tables_and_host_resource_data(instance)
+            .0
+    }
+
+    pub(crate) fn component_resource_tables_and_host_resource_data(
+        &mut self,
+        instance: Option<Instance>,
+    ) -> (
+        vm::component::ResourceTables<'_>,
+        &mut crate::component::HostResourceData,
+    ) {
+        let store_id = self.id();
+        let data = self.component_data_mut();
+        let guest = instance.map(|i| {
+            let i = i.id();
+            i.assert_belongs_to(store_id);
+            data.instances[i.instance]
+                .as_mut()
+                .unwrap()
+                .get_mut()
+                .instance_states()
+        });
+
+        (
+            vm::component::ResourceTables {
+                host_table: Some(&mut data.component_host_table),
+                calls: &mut data.component_calls,
+                guest,
+            },
+            &mut data.host_resource_data,
+        )
+    }
+}

crates/wasmtime/src/runtime/instance.rs

@@ -225,7 +225,8 @@ impl Instance {
         // Note that under normal operation this shouldn't do much as the list
         // of funcs-with-holes should generally be empty. As a result the
         // process of filling this out is not super optimized at this point.
-        store.register_module(module)?;
+        let (modules, engine) = store.modules_and_engine_mut();
+        modules.register_module(module, engine)?;
         let (funcrefs, modules) = store.func_refs_and_modules();
         funcrefs.fill(modules);
 
@@ -304,7 +305,8 @@ impl Instance {
 
         // Register the module just before instantiation to ensure we keep the module
         // properly referenced while in use by the store.
-        let module_id = store.register_module(module)?;
+        let (modules, engine) = store.modules_and_engine_mut();
+        let module_id = modules.register_module(module, engine)?;
 
         // The first thing we do is issue an instance allocation request
         // to the instance allocator. This, on success, will give us an
@@ -950,7 +952,8 @@ fn pre_instantiate_raw(
 ) -> Result<OwnedImports> {
     // Register this module and use it to fill out any funcref wasm_call holes
     // we can. For more comments on this see `typecheck_externs`.
-    store.register_module(module)?;
+    let (modules, engine) = store.modules_and_engine_mut();
+    modules.register_module(module, engine)?;
     let (funcrefs, modules) = store.func_refs_and_modules();
     funcrefs.fill(modules);