Add OOM test for component ResourceAny::resource_drop_async (#13051)

Convert several infallible allocations to fallible alternatives:
- HandleTable::slots: Vec<Slot> -> TryVec<Slot>
- HostResourceData::table_slot_metadata: Vec<TableSlot> -> TryVec<TableSlot>
- Resource type tracking: PrimaryMap<ResourceIndex, ResourceType> ->
  TryPrimaryMap across vm/component.rs, instance.rs, matching.rs,
  types.rs, component.rs, and linker.rs
- Backtrace trace_through_continuations: remove Vec::collect() calls,
  iterate directly to avoid allocation
- catch_traps: Box::new(Trap) -> try_new::<Box<_>>(Trap)

Author: fitzgen

crates/fuzzing/tests/oom/component_resource.rs

@@ -0,0 +1,48 @@
+#![cfg(arc_try_new)]
+
+use wasmtime::component::{Component, Linker, ResourceAny};
+use wasmtime::{Config, Engine, Result, Store};
+use wasmtime_fuzzing::oom::OomTest;
+
+#[tokio::test]
+async fn component_resource_any_resource_drop_async() -> Result<()> {
+    let component_bytes = {
+        let mut config = Config::new();
+        config.concurrency_support(false);
+        let engine = Engine::new(&config)?;
+        Component::new(
+            &engine,
+            r#"
+                (component
+                    (type $t' (resource (rep i32)))
+                    (export $t "t" (type $t'))
+
+                    (core func $new (canon resource.new $t))
+                    (func (export "mk") (param "r" u32) (result (own $t))
+                        (canon lift (core func $new))
+                    )
+                )
+            "#,
+        )?
+        .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()
+        .allow_alloc_after_oom(true)
+        .test_async(|| async {
+            let mut store = Store::try_new(&engine, ())?;
+            let instance = instance_pre.instantiate_async(&mut store).await?;
+            let mk = instance.get_typed_func::<(u32,), (ResourceAny,)>(&mut store, "mk")?;
+            let (resource,) = mk.call_async(&mut store, (42,)).await?;
+            resource.resource_drop_async(&mut store).await?;
+            Ok(())
+        })
+        .await
+}

crates/fuzzing/tests/oom/main.rs

@@ -7,6 +7,7 @@ mod btree_map;
 mod caller;
 mod component_func;
 mod component_linker;
+mod component_resource;
 mod config;
 mod engine;
 mod entity_set;

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

@@ -388,7 +388,7 @@ impl Component {
     }
 
     fn with_uninstantiated_instance_type<R>(&self, f: impl FnOnce(&InstanceType<'_>) -> R) -> R {
-        let resources = Arc::new(PrimaryMap::new());
+        let resources = Arc::new(TryPrimaryMap::new());
         f(&InstanceType {
             types: self.types(),
             resources: &resources,

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

@@ -715,7 +715,7 @@ pub(crate) enum RuntimeImport {
     },
 }
 
-pub type ImportedResources = PrimaryMap<ResourceIndex, ResourceType>;
+pub type ImportedResources = TryPrimaryMap<ResourceIndex, ResourceType>;
 
 impl<'a> Instantiator<'a> {
     fn new(
@@ -727,7 +727,7 @@ impl<'a> Instantiator<'a> {
         let (modules, engine, breakpoints) = store.modules_and_engine_and_breakpoints_mut();
         modules.register_component(component, engine, breakpoints)?;
         let imported_resources: ImportedResources =
-            PrimaryMap::with_capacity(env_component.imported_resources.len());
+            TryPrimaryMap::with_capacity(env_component.imported_resources.len())?;
 
         let instance = ComponentInstance::new(
             store.store_data().components.next_component_instance_id(),
@@ -763,7 +763,7 @@ impl<'a> Instantiator<'a> {
                 } => (*ty, NonNull::from(dtor_funcref)),
                 _ => unreachable!(),
             };
-            let i = self.instance_resource_types_mut(store.0).push(ty);
+            let i = self.instance_resource_types_mut(store.0).push(ty)?;
             assert_eq!(i, idx);
             self.instance_mut(store.0)
                 .set_resource_destructor(idx, Some(func_ref));
@@ -906,13 +906,13 @@ impl<'a> Instantiator<'a> {
                     self.extract_post_return(store.0, post_return)
                 }
 
-                GlobalInitializer::Resource(r) => self.resource(store.0, r),
+                GlobalInitializer::Resource(r) => self.resource(store.0, r)?,
             }
         }
         Ok(())
     }
 
-    fn resource(&mut self, store: &mut StoreOpaque, resource: &Resource) {
+    fn resource(&mut self, store: &mut StoreOpaque, resource: &Resource) -> Result<()> {
         let dtor = resource
             .dtor
             .as_ref()
@@ -929,8 +929,9 @@ impl<'a> Instantiator<'a> {
         let ty = ResourceType::guest(store.id(), instance, resource.index);
         self.instance_mut(store)
             .set_resource_destructor(index, dtor);
-        let i = self.instance_resource_types_mut(store).push(ty);
+        let i = self.instance_resource_types_mut(store).push(ty)?;
         debug_assert_eq!(i, index);
+        Ok(())
     }
 
     fn extract_memory(&mut self, store: &mut StoreOpaque, memory: &ExtractMemory) {
@@ -1088,7 +1089,7 @@ impl<'a> Instantiator<'a> {
 pub struct InstancePre<T: 'static> {
     component: Component,
     imports: Arc<PrimaryMap<RuntimeImportIndex, RuntimeImport>>,
-    resource_types: Arc<PrimaryMap<ResourceIndex, ResourceType>>,
+    resource_types: Arc<TryPrimaryMap<ResourceIndex, ResourceType>>,
     asyncness: Asyncness,
     _marker: marker::PhantomData<fn() -> T>,
 }
@@ -1116,7 +1117,7 @@ impl<T: 'static> InstancePre<T> {
     pub(crate) unsafe fn new_unchecked(
         component: Component,
         imports: Arc<PrimaryMap<RuntimeImportIndex, RuntimeImport>>,
-        resource_types: Arc<PrimaryMap<ResourceIndex, ResourceType>>,
+        resource_types: Arc<TryPrimaryMap<ResourceIndex, ResourceType>>,
     ) -> InstancePre<T> {
         let mut asyncness = Asyncness::No;
         for (_, import) in imports.iter() {

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

@@ -177,7 +177,7 @@ impl<T: 'static> Linker<T> {
             engine: &self.engine,
             types: component.types(),
             strings: &self.strings,
-            imported_resources: try_new::<Arc<_>>(Default::default())?,
+            imported_resources: try_new::<Arc<_>>(TryPrimaryMap::new())?,
         };
 
         // Walk over the component's list of import names and use that to lookup

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

@@ -7,25 +7,25 @@ use crate::runtime::vm::component::ComponentInstance;
 use crate::types::matching;
 use crate::{Engine, prelude::*};
 use alloc::sync::Arc;
-use wasmtime_environ::PrimaryMap;
 use wasmtime_environ::component::{
     ComponentTypes, NameMap, ResourceIndex, TypeComponentInstance, TypeDef, TypeFuncIndex,
     TypeFutureTableIndex, TypeModule, TypeResourceTable, TypeResourceTableIndex,
     TypeStreamTableIndex,
 };
+use wasmtime_environ::prelude::TryPrimaryMap;
 
 pub struct TypeChecker<'a> {
     pub engine: &'a Engine,
     pub types: &'a Arc<ComponentTypes>,
     pub strings: &'a Strings,
-    pub imported_resources: Arc<PrimaryMap<ResourceIndex, ResourceType>>,
+    pub imported_resources: Arc<TryPrimaryMap<ResourceIndex, ResourceType>>,
 }
 
 #[derive(Copy, Clone)]
 #[doc(hidden)]
 pub struct InstanceType<'a> {
     pub types: &'a Arc<ComponentTypes>,
-    pub resources: &'a Arc<PrimaryMap<ResourceIndex, ResourceType>>,
+    pub resources: &'a Arc<TryPrimaryMap<ResourceIndex, ResourceType>>,
 }
 
 impl TypeChecker<'_> {
@@ -89,7 +89,7 @@ impl TypeChecker<'_> {
                     // cloned.
                     None => {
                         let resources = Arc::get_mut(&mut self.imported_resources).unwrap();
-                        let id = resources.push(*actual);
+                        let id = resources.push(*actual)?;
                         assert_eq!(id, i);
                     }
 

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

@@ -62,7 +62,7 @@ pub struct HostResourceTables<'a> {
 #[derive(Default)]
 pub struct HostResourceData {
     cur_generation: u32,
-    table_slot_metadata: Vec<TableSlot>,
+    table_slot_metadata: TryVec<TableSlot>,
 }
 
 #[derive(Copy, Clone)]
@@ -145,15 +145,15 @@ impl<'a> HostResourceTables<'a> {
         instance: Option<RuntimeInstance>,
     ) -> Result<HostResourceIndex> {
         let idx = self.tables.resource_lower_own(TypedResource::Host(rep))?;
-        Ok(self.new_host_index(idx, dtor, instance))
+        Ok(self.new_host_index(idx, dtor, instance)?)
     }
 
     /// See [`HostResourceTables::host_resource_lower_own`].
     pub fn host_resource_lower_borrow(&mut self, rep: u32) -> Result<HostResourceIndex> {
         let idx = self
             .tables
             .resource_lower_borrow(TypedResource::Host(rep))?;
-        Ok(self.new_host_index(idx, None, None))
+        Ok(self.new_host_index(idx, None, None)?)
     }
 
     /// Validates that `idx` is still valid for the host tables, notably
@@ -201,7 +201,7 @@ impl<'a> HostResourceTables<'a> {
         idx: u32,
         dtor: Option<NonNull<VMFuncRef>>,
         instance: Option<RuntimeInstance>,
-    ) -> HostResourceIndex {
+    ) -> Result<HostResourceIndex> {
         let list = &mut self.host_resource_data.table_slot_metadata;
         let info = TableSlot {
             generation: self.host_resource_data.cur_generation,
@@ -219,14 +219,14 @@ impl<'a> HostResourceTables<'a> {
                         generation: 0,
                         instance: None,
                         dtor: None,
-                    });
+                    })?;
                 }
                 assert_eq!(idx as usize, list.len());
-                list.push(info);
+                list.push(info)?;
             }
         }
 
-        HostResourceIndex::new(idx, info.generation)
+        Ok(HostResourceIndex::new(idx, info.generation))
     }
 
     /// Drops a host-owned resource from host tables.

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

@@ -7,14 +7,14 @@ use crate::{Engine, ExternType, FuncType, prelude::*};
 use alloc::sync::Arc;
 use core::fmt;
 use core::ops::Deref;
+use wasmtime_environ::PanicOnOom as _;
 use wasmtime_environ::component::{
     ComponentTypes, Export, InterfaceType, ResourceIndex, TypeComponentIndex,
     TypeComponentInstanceIndex, TypeDef, TypeEnumIndex, TypeFlagsIndex, TypeFuncIndex,
     TypeFutureIndex, TypeFutureTableIndex, TypeListIndex, TypeMapIndex, TypeModuleIndex,
     TypeOptionIndex, TypeRecordIndex, TypeResourceTable, TypeResourceTableIndex, TypeResultIndex,
     TypeStreamIndex, TypeStreamTableIndex, TypeTupleIndex, TypeVariantIndex,
 };
-use wasmtime_environ::{PanicOnOom as _, PrimaryMap};
 
 pub use crate::component::resources::ResourceType;
 
@@ -40,7 +40,7 @@ pub use crate::component::resources::ResourceType;
 struct Handle<T> {
     index: T,
     types: Arc<ComponentTypes>,
-    resources: Arc<PrimaryMap<ResourceIndex, ResourceType>>,
+    resources: Arc<TryPrimaryMap<ResourceIndex, ResourceType>>,
 }
 
 impl<T> Handle<T> {
@@ -94,9 +94,9 @@ impl<T: fmt::Debug> fmt::Debug for Handle<T> {
 /// Type checker between two `Handle`s
 struct TypeChecker<'a> {
     a_types: &'a ComponentTypes,
-    a_resource: &'a PrimaryMap<ResourceIndex, ResourceType>,
+    a_resource: &'a TryPrimaryMap<ResourceIndex, ResourceType>,
     b_types: &'a ComponentTypes,
-    b_resource: &'a PrimaryMap<ResourceIndex, ResourceType>,
+    b_resource: &'a TryPrimaryMap<ResourceIndex, ResourceType>,
 }
 
 impl TypeChecker<'_> {

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

@@ -138,7 +138,7 @@ pub struct ComponentInstance {
 
     /// Storage for the type information about resources within this component
     /// instance.
-    resource_types: Arc<PrimaryMap<ResourceIndex, ResourceType>>,
+    resource_types: Arc<TryPrimaryMap<ResourceIndex, ResourceType>>,
 
     /// Arguments that this instance used to be instantiated.
     ///
@@ -319,7 +319,7 @@ impl ComponentInstance {
     pub(crate) fn new(
         id: ComponentInstanceId,
         component: &Component,
-        resource_types: Arc<PrimaryMap<ResourceIndex, ResourceType>>,
+        resource_types: Arc<TryPrimaryMap<ResourceIndex, ResourceType>>,
         imports: &Arc<PrimaryMap<RuntimeImportIndex, RuntimeImport>>,
         store: NonNull<dyn VMStore>,
     ) -> Result<OwnedComponentInstance, OutOfMemory> {
@@ -826,14 +826,14 @@ impl ComponentInstance {
     }
 
     /// Returns a reference to the resource type information.
-    pub fn resource_types(&self) -> &Arc<PrimaryMap<ResourceIndex, ResourceType>> {
+    pub fn resource_types(&self) -> &Arc<TryPrimaryMap<ResourceIndex, ResourceType>> {
         &self.resource_types
     }
 
     /// Returns a mutable reference to the resource type information.
     pub fn resource_types_mut(
         self: Pin<&mut Self>,
-    ) -> &mut Arc<PrimaryMap<ResourceIndex, ResourceType>> {
+    ) -> &mut Arc<TryPrimaryMap<ResourceIndex, ResourceType>> {
         // SAFETY: we've chosen the `Pin` guarantee of `Self` to not apply to
         // the map returned.
         unsafe { &mut self.get_unchecked_mut().resource_types }

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

@@ -1,6 +1,6 @@
 use super::{TypedResource, TypedResourceIndex};
+use crate::prelude::TryVec;
 use crate::{Result, bail};
-use alloc::vec::Vec;
 use core::mem;
 use wasmtime_environ::component::{TypeFutureTableIndex, TypeStreamTableIndex};
 
@@ -114,14 +114,14 @@ enum Slot {
 
 pub struct HandleTable {
     next: u32,
-    slots: Vec<Slot>,
+    slots: TryVec<Slot>,
 }
 
 impl Default for HandleTable {
     fn default() -> Self {
         Self {
             next: 0,
-            slots: Vec::new(),
+            slots: TryVec::new(),
         }
     }
 }
@@ -139,7 +139,7 @@ impl HandleTable {
         if next == self.slots.len() {
             self.slots.push(Slot::Free {
                 next: self.next.checked_add(1).unwrap(),
-            });
+            })?;
         }
         let ret = self.next;
         self.next = match mem::replace(&mut self.slots[next], slot) {