Add OOM test for component InstancePre::instantiate_async (#12993)

* Add OOM test for component InstancePre::instantiate_async

Fix infallible allocations in the component model instantiation path:

- Use TryPrimaryMap for instance_states and instances in ComponentInstance
- Use TryPrimaryMap for instances in ComponentStoreData
- Use TryVec for scopes in ComponentTasksNotConcurrent
- Use TryPrimaryMap for guest field in ResourceTables
- Use try_new for Arc and Box allocations in instantiation/call paths
- Make push_component_instance and push_instance_id return Result
- Make enter_call_not_concurrent return Result

* remove now-unused import

* fix imports for different cfgs

* really fix imports this time

Author: fitzgen

crates/fuzzing/tests/oom/component_func.rs

@@ -0,0 +1,44 @@
+#![cfg(arc_try_new)]
+
+use wasmtime::component::{Component, Linker};
+use wasmtime::{Config, Engine, Result, Store};
+use wasmtime_fuzzing::oom::OomTest;
+
+#[tokio::test]
+async fn component_instance_pre_instantiate_async() -> Result<()> {
+    let component_bytes = {
+        let mut config = Config::new();
+        config.concurrency_support(false);
+        let engine = Engine::new(&config)?;
+        Component::new(
+            &engine,
+            r#"
+                (component
+                    (core module $m
+                        (func (export "id") (param i32) (result i32) (local.get 0))
+                    )
+                    (core instance $i (instantiate $m))
+                    (func (export "id") (param "x" s32) (result s32)
+                        (canon lift (core func $i "id"))
+                    )
+                )
+            "#,
+        )?
+        .serialize()?
+    };
+    let mut config = Config::new();
+    config.enable_compiler(false);
+    config.concurrency_support(false);
+    let engine = Engine::new(&config)?;
+    let component = unsafe { Component::deserialize(&engine, &component_bytes)? };
+    let linker = Linker::<()>::new(&engine);
+    let instance_pre = linker.instantiate_pre(&component)?;
+
+    OomTest::new()
+        .test_async(|| async {
+            let mut store = Store::try_new(&engine, ())?;
+            let _instance = instance_pre.instantiate_async(&mut store).await?;
+            Ok(())
+        })
+        .await
+}

crates/fuzzing/tests/oom/main.rs

@@ -5,6 +5,7 @@ mod bit_set;
 mod boxed;
 mod btree_map;
 mod caller;
+mod component_func;
 mod config;
 mod engine;
 mod entity_set;

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

@@ -1450,7 +1450,7 @@ impl StoreOpaque {
     ) -> Result<()> {
         log::trace!("enter sync call {callee:?}");
         if !self.concurrency_support() {
-            return Ok(self.enter_call_not_concurrent());
+            return self.enter_call_not_concurrent();
         }
 
         let state = self.concurrent_state_mut();
@@ -1548,7 +1548,7 @@ impl StoreOpaque {
     /// situations.
     pub(crate) fn host_task_create(&mut self) -> Result<Option<TableId<HostTask>>> {
         if !self.concurrency_support() {
-            self.enter_call_not_concurrent();
+            self.enter_call_not_concurrent()?;
             return Ok(None);
         }
         let state = self.concurrent_state_mut();

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

@@ -168,15 +168,15 @@ impl StoreOpaque {
         _callee_async: bool,
         _callee: RuntimeInstance,
     ) -> Result<()> {
-        Ok(self.enter_call_not_concurrent())
+        self.enter_call_not_concurrent()
     }
 
     pub(crate) fn exit_guest_sync_call(&mut self) -> Result<()> {
         Ok(self.exit_call_not_concurrent())
     }
 
     pub(crate) fn host_task_create(&mut self) -> Result<()> {
-        Ok(self.enter_call_not_concurrent())
+        self.enter_call_not_concurrent()
     }
 
     pub(crate) fn host_task_reenter_caller(&mut self) -> Result<()> {

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

@@ -771,15 +771,15 @@ impl Func {
         };
         if results_ty.abi.flat_count(max_flat).is_some() {
             let mut flat = src.iter();
-            Ok(Box::new(
+            Ok(try_new::<Box<_>>(
                 results_ty
                     .types
                     .iter()
                     .map(move |ty| Val::lift(cx, *ty, &mut flat)),
-            ))
+            )?)
         } else {
             let iter = Self::load_results(cx, results_ty, &mut src.iter())?;
-            Ok(Box::new(iter))
+            Ok(try_new::<Box<_>>(iter)?)
         }
     }
 

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

@@ -732,11 +732,11 @@ impl<'a> Instantiator<'a> {
         let instance = ComponentInstance::new(
             store.store_data().components.next_component_instance_id(),
             component,
-            Arc::new(imported_resources),
+            try_new::<Arc<_>>(imported_resources)?,
             imports,
             store.traitobj(),
         )?;
-        let id = store.store_data_mut().push_component_instance(instance);
+        let id = store.store_data_mut().push_component_instance(instance)?;
 
         Ok(Instantiator {
             component,
@@ -878,7 +878,7 @@ impl<'a> Instantiator<'a> {
                         store.0.exit_guest_sync_call()?;
                     }
 
-                    self.instance_mut(store.0).push_instance_id(i.id());
+                    self.instance_mut(store.0).push_instance_id(i.id())?;
                 }
 
                 GlobalInitializer::LowerImport { import, index } => {

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

@@ -1,19 +1,21 @@
 use crate::prelude::*;
-#[cfg(feature = "component-model-async")]
-use crate::runtime::component::ResourceTable;
 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::{
     CallContext, ComponentInstance, HandleTable, OwnedComponentInstance,
 };
 use crate::store::{StoreData, StoreId, StoreOpaque};
 use crate::{AsContext, AsContextMut, Engine, Store, StoreContextMut};
 use core::pin::Pin;
-use wasmtime_environ::PrimaryMap;
 use wasmtime_environ::component::RuntimeComponentInstanceIndex;
+use wasmtime_environ::prelude::TryPrimaryMap;
+
+#[cfg(feature = "component-model-async")]
+use crate::{
+    component::ResourceTable,
+    runtime::vm::{VMStore, component::InstanceState},
+};
 
 /// Default amount of fuel allowed for all guest-to-host calls in the component
 /// model.
@@ -29,7 +31,7 @@ const DEFAULT_HOSTCALL_FUEL: usize = 128 << 20;
 pub struct ComponentStoreData {
     /// All component instances, in a similar manner to how core wasm instances
     /// are managed.
-    instances: PrimaryMap<ComponentInstanceId, Option<OwnedComponentInstance>>,
+    instances: TryPrimaryMap<ComponentInstanceId, Option<OwnedComponentInstance>>,
 
     /// Whether an instance belonging to this store has trapped.
     trapped: bool,
@@ -151,15 +153,10 @@ impl ComponentStoreData {
                 continue;
             };
 
-            assert!(
-                instance
-                    .get_mut()
-                    .instance_states()
-                    .0
-                    .iter_mut()
-                    .all(|(_, state)| state.handle_table().is_empty()
-                        && state.concurrent_state().pending_is_empty())
-            );
+            assert!(instance.get_mut().instance_states().0.iter_mut().all(
+                |(_, state): (_, &mut InstanceState)| state.handle_table().is_empty()
+                    && state.concurrent_state().pending_is_empty()
+            ));
         }
     }
 
@@ -259,11 +256,11 @@ impl StoreData {
     pub(crate) fn push_component_instance(
         &mut self,
         data: OwnedComponentInstance,
-    ) -> ComponentInstanceId {
+    ) -> Result<ComponentInstanceId, OutOfMemory> {
         let expected = data.get().id();
-        let ret = self.components.instances.push(Some(data));
+        let ret = self.components.instances.push(Some(data))?;
         assert_eq!(expected, ret);
-        ret
+        Ok(ret)
     }
 
     pub(crate) fn component_instance(&self, id: ComponentInstanceId) -> &ComponentInstance {
@@ -394,12 +391,13 @@ impl StoreOpaque {
         )
     }
 
-    pub(crate) fn enter_call_not_concurrent(&mut self) {
+    pub(crate) fn enter_call_not_concurrent(&mut self) -> Result<()> {
         let state = match &mut self.component_data_mut().task_state {
             ComponentTaskState::NotConcurrent(state) => state,
             ComponentTaskState::Concurrent(_) => unreachable!(),
         };
-        state.scopes.push(CallContext::default());
+        state.scopes.push(CallContext::default())?;
+        Ok(())
     }
 
     pub(crate) fn exit_call_not_concurrent(&mut self) {
@@ -499,7 +497,7 @@ impl<T> StoreContextMut<'_, T> {
 
 #[derive(Default)]
 pub struct ComponentTasksNotConcurrent {
-    scopes: Vec<CallContext>,
+    scopes: TryVec<CallContext>,
 }
 
 impl ComponentTaskState {

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

@@ -29,6 +29,7 @@ use core::pin::Pin;
 use core::ptr::NonNull;
 use wasmtime_environ::component::*;
 use wasmtime_environ::error::OutOfMemory;
+use wasmtime_environ::prelude::TryPrimaryMap;
 use wasmtime_environ::{HostPtr, PrimaryMap, VMSharedTypeIndex};
 
 #[allow(
@@ -129,11 +130,11 @@ pub struct ComponentInstance {
 
     /// Contains state specific to each (sub-)component instance within this
     /// top-level instance.
-    instance_states: PrimaryMap<RuntimeComponentInstanceIndex, InstanceState>,
+    instance_states: TryPrimaryMap<RuntimeComponentInstanceIndex, InstanceState>,
 
     /// What all compile-time-identified core instances are mapped to within the
     /// `Store` that this component belongs to.
-    instances: PrimaryMap<RuntimeInstanceIndex, InstanceId>,
+    instances: TryPrimaryMap<RuntimeInstanceIndex, InstanceId>,
 
     /// Storage for the type information about resources within this component
     /// instance.
@@ -324,22 +325,24 @@ impl ComponentInstance {
     ) -> Result<OwnedComponentInstance, OutOfMemory> {
         let offsets = VMComponentOffsets::new(HostPtr, component.env_component());
         let num_instances = component.env_component().num_runtime_component_instances;
-        let mut instance_states = PrimaryMap::with_capacity(num_instances.try_into().unwrap());
+        let mut instance_states = TryPrimaryMap::with_capacity(num_instances.try_into().unwrap())?;
         for _ in 0..num_instances {
-            instance_states.push(InstanceState::default());
+            instance_states.push(InstanceState::default())?;
         }
 
+        let instances = TryPrimaryMap::with_capacity(
+            component
+                .env_component()
+                .num_runtime_instances
+                .try_into()
+                .unwrap(),
+        )?;
+
         let mut ret = OwnedInstance::new(ComponentInstance {
             id,
             offsets,
             instance_states,
-            instances: PrimaryMap::with_capacity(
-                component
-                    .env_component()
-                    .num_runtime_instances
-                    .try_into()
-                    .unwrap(),
-            ),
+            instances,
             component: component.clone(),
             resource_types,
             imports: imports.clone(),
@@ -860,7 +863,7 @@ impl ComponentInstance {
     pub fn instance_states(
         self: Pin<&mut Self>,
     ) -> (
-        &mut PrimaryMap<RuntimeComponentInstanceIndex, InstanceState>,
+        &mut TryPrimaryMap<RuntimeComponentInstanceIndex, InstanceState>,
         &ComponentTypes,
     ) {
         // safety: we've chosen the `pin` guarantee of `self` to not apply to
@@ -906,7 +909,10 @@ impl ComponentInstance {
 
     /// Pushes a new runtime instance that's been created into
     /// `self.instances`.
-    pub fn push_instance_id(self: Pin<&mut Self>, id: InstanceId) -> RuntimeInstanceIndex {
+    pub fn push_instance_id(
+        self: Pin<&mut Self>,
+        id: InstanceId,
+    ) -> Result<RuntimeInstanceIndex, OutOfMemory> {
         self.instances_mut().push(id)
     }
 
@@ -920,7 +926,7 @@ impl ComponentInstance {
         self.instances[idx]
     }
 
-    fn instances_mut(self: Pin<&mut Self>) -> &mut PrimaryMap<RuntimeInstanceIndex, InstanceId> {
+    fn instances_mut(self: Pin<&mut Self>) -> &mut TryPrimaryMap<RuntimeInstanceIndex, InstanceId> {
         // SAFETY: we've chosen the `Pin` guarantee of `Self` to not apply to
         // the map returned.
         unsafe { &mut self.get_unchecked_mut().instances }

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

@@ -26,10 +26,10 @@ use crate::prelude::*;
 use core::error::Error;
 use core::fmt;
 use core::mem;
-use wasmtime_environ::PrimaryMap;
 use wasmtime_environ::component::{
     ComponentTypes, RuntimeComponentInstanceIndex, TypeResourceTableIndex,
 };
+use wasmtime_environ::prelude::TryPrimaryMap;
 
 /// Contextual state necessary to perform resource-related operations.
 ///
@@ -55,7 +55,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, InstanceState>,
+        &'a mut TryPrimaryMap<RuntimeComponentInstanceIndex, InstanceState>,
         &'a ComponentTypes,
     )>,