Make component linker allocations fallible (#13083)

* Make component linker allocations fallible

* Address review feedback

Author: fitzgen

crates/environ/src/component/names.rs

@@ -1,6 +1,6 @@
 use crate::collections::TryCow;
 use crate::error::{Result, bail};
-use crate::prelude::*;
+use crate::{Atom, StringPool, prelude::*};
 use alloc::sync::Arc;
 use core::borrow::Borrow;
 use core::hash::Hash;
@@ -263,6 +263,19 @@ impl NameMapIntern for NameMapNoIntern {
     }
 }
 
+impl NameMapIntern for StringPool {
+    type Key = Atom;
+    type BorrowedKey = Atom;
+
+    fn intern(&mut self, string: &str) -> Result<Atom, OutOfMemory> {
+        self.insert(string)
+    }
+
+    fn lookup(&self, string: &str) -> Option<TryCow<'_, Atom>> {
+        self.get_atom(string).map(TryCow::Owned)
+    }
+}
+
 /// Determines a version-based "alternate lookup key" for the `name` specified.
 ///
 /// Some examples are:

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

@@ -76,22 +76,22 @@ impl HostFunc {
     ///
     /// Note that if `asyncness` is mistaken then that'll result in panics
     /// in Wasmtime, but not memory unsafety.
-    fn new<T, F, P, R>(asyncness: Asyncness, func: F) -> Arc<HostFunc>
+    fn new<T, F, P, R>(asyncness: Asyncness, func: F) -> Result<Arc<HostFunc>>
     where
         T: 'static,
         R: Send + Sync + 'static,
         F: HostFn<T, P, R> + Send + Sync + 'static,
     {
-        Arc::new(HostFunc {
+        Ok(try_new::<Arc<_>>(HostFunc {
             entrypoint: F::cabi_entrypoint,
             typecheck: F::typecheck,
-            func: Box::new(func),
+            func: try_new::<Box<_>>(func)?,
             asyncness,
-        })
+        })?)
     }
 
     /// Equivalent for `Linker::func_wrap`
-    pub(crate) fn func_wrap<T, F, P, R>(func: F) -> Arc<HostFunc>
+    pub(crate) fn func_wrap<T, F, P, R>(func: F) -> Result<Arc<HostFunc>>
     where
         T: 'static,
         F: Fn(StoreContextMut<T>, P) -> Result<R> + Send + Sync + 'static,
@@ -108,7 +108,7 @@ impl HostFunc {
 
     /// Equivalent for `Linker::func_wrap_async`
     #[cfg(feature = "async")]
-    pub(crate) fn func_wrap_async<T, F, P, R>(func: F) -> Arc<HostFunc>
+    pub(crate) fn func_wrap_async<T, F, P, R>(func: F) -> Result<Arc<HostFunc>>
     where
         T: 'static,
         F: Fn(StoreContextMut<'_, T>, P) -> Box<dyn Future<Output = Result<R>> + Send + '_>
@@ -132,7 +132,7 @@ impl HostFunc {
 
     /// Equivalent for `Linker::func_wrap_concurrent`
     #[cfg(feature = "component-model-async")]
-    pub(crate) fn func_wrap_concurrent<T, F, P, R>(func: F) -> Arc<HostFunc>
+    pub(crate) fn func_wrap_concurrent<T, F, P, R>(func: F) -> Result<Arc<HostFunc>>
     where
         T: 'static,
         F: Fn(&Accessor<T>, P) -> Pin<Box<dyn Future<Output = Result<R>> + Send + '_>>
@@ -155,7 +155,7 @@ impl HostFunc {
     }
 
     /// Equivalent of `Linker::func_new`
-    pub(crate) fn func_new<T, F>(func: F) -> Arc<HostFunc>
+    pub(crate) fn func_new<T, F>(func: F) -> Result<Arc<HostFunc>>
     where
         T: 'static,
         F: Fn(StoreContextMut<'_, T>, ComponentFunc, &[Val], &mut [Val]) -> Result<()>
@@ -177,7 +177,7 @@ impl HostFunc {
 
     /// Equivalent of `Linker::func_new_async`
     #[cfg(feature = "async")]
-    pub(crate) fn func_new_async<T, F>(func: F) -> Arc<HostFunc>
+    pub(crate) fn func_new_async<T, F>(func: F) -> Result<Arc<HostFunc>>
     where
         T: 'static,
         F: for<'a> Fn(
@@ -209,7 +209,7 @@ impl HostFunc {
 
     /// Equivalent of `Linker::func_new_concurrent`
     #[cfg(feature = "component-model-async")]
-    pub(crate) fn func_new_concurrent<T, F>(func: F) -> Arc<HostFunc>
+    pub(crate) fn func_new_concurrent<T, F>(func: F) -> Result<Arc<HostFunc>>
     where
         T: 'static,
         F: for<'a> Fn(

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

@@ -7,16 +7,14 @@ use crate::component::types;
 use crate::component::{
     Component, ComponentNamedList, Instance, InstancePre, Lift, Lower, ResourceType, Val,
 };
-use crate::hash_map::HashMap;
 use crate::prelude::*;
 use crate::{AsContextMut, Engine, Module, StoreContextMut};
 use alloc::sync::Arc;
 use core::marker;
 #[cfg(feature = "component-model-async")]
 use core::pin::Pin;
-use wasmtime_environ::PrimaryMap;
-use wasmtime_environ::collections::TryCow;
-use wasmtime_environ::component::{NameMap, NameMapIntern};
+use wasmtime_environ::component::NameMap;
+use wasmtime_environ::{Atom, PrimaryMap, StringPool};
 
 /// A type used to instantiate [`Component`]s.
 ///
@@ -61,9 +59,9 @@ use wasmtime_environ::component::{NameMap, NameMapIntern};
 /// be instantiated.
 pub struct Linker<T: 'static> {
     engine: Engine,
-    strings: Strings,
-    map: NameMap<usize, Definition>,
-    path: Vec<usize>,
+    strings: StringPool,
+    map: NameMap<Atom, Definition>,
+    path: Vec<Atom>,
     allow_shadowing: bool,
     _marker: marker::PhantomData<fn() -> T>,
 }
@@ -72,7 +70,7 @@ impl<T: 'static> Clone for Linker<T> {
     fn clone(&self) -> Linker<T> {
         Linker {
             engine: self.engine.clone(),
-            strings: self.strings.clone(),
+            strings: self.strings.clone_panic_on_oom(),
             map: self.map.clone_panic_on_oom(),
             path: self.path.clone(),
             allow_shadowing: self.allow_shadowing,
@@ -81,30 +79,24 @@ impl<T: 'static> Clone for Linker<T> {
     }
 }
 
-#[derive(Clone, Default)]
-pub struct Strings {
-    string2idx: HashMap<Arc<str>, usize>,
-    strings: Vec<Arc<str>>,
-}
-
 /// Structure representing an "instance" being defined within a linker.
 ///
 /// Instances do not need to be actual [`Instance`]s and instead are defined by
 /// a "bag of named items", so each [`LinkerInstance`] can further define items
 /// internally.
 pub struct LinkerInstance<'a, T: 'static> {
     engine: &'a Engine,
-    path: &'a mut Vec<usize>,
+    path: &'a mut Vec<Atom>,
     path_len: usize,
-    strings: &'a mut Strings,
-    map: &'a mut NameMap<usize, Definition>,
+    strings: &'a mut StringPool,
+    map: &'a mut NameMap<Atom, Definition>,
     allow_shadowing: bool,
     _marker: marker::PhantomData<fn() -> T>,
 }
 
 #[derive(Debug)]
 pub(crate) enum Definition {
-    Instance(NameMap<usize, Definition>),
+    Instance(NameMap<Atom, Definition>),
     Func(Arc<HostFunc>),
     Module(Module),
     Resource(ResourceType, Arc<crate::func::HostFunc>),
@@ -127,7 +119,7 @@ impl<T: 'static> Linker<T> {
     pub fn new(engine: &Engine) -> Linker<T> {
         Linker {
             engine: engine.clone(),
-            strings: Strings::default(),
+            strings: StringPool::default(),
             map: NameMap::default(),
             allow_shadowing: false,
             path: Vec::new(),
@@ -345,9 +337,20 @@ impl<T: 'static> Linker<T> {
 
             match item_def {
                 TypeDef::ComponentFunc(_) => {
-                    let fully_qualified_name = parent_instance
-                        .map(|parent| format!("{parent}#{item_name}"))
-                        .unwrap_or_else(|| item_name.to_owned());
+                    let fully_qualified_name = match parent_instance {
+                        Some(parent) => {
+                            let mut s = TryString::new();
+                            s.push_str(parent)?;
+                            s.push('#')?;
+                            s.push_str(item_name)?;
+                            s
+                        }
+                        None => {
+                            let mut s = TryString::new();
+                            s.push_str(item_name)?;
+                            s
+                        }
+                    };
                     linker.func_new(&item_name, move |_, _, _, _| {
                         bail!("unknown import: `{fully_qualified_name}` has not been defined")
                     })?;
@@ -438,7 +441,7 @@ impl<T: 'static> LinkerInstance<'_, T> {
         Params: ComponentNamedList + Lift + 'static,
         Return: ComponentNamedList + Lower + 'static,
     {
-        self.insert(name, Definition::Func(HostFunc::func_wrap(func)))?;
+        self.insert(name, Definition::Func(HostFunc::func_wrap(func)?))?;
         Ok(())
     }
 
@@ -485,7 +488,7 @@ impl<T: 'static> LinkerInstance<'_, T> {
         Params: ComponentNamedList + Lift + 'static,
         Return: ComponentNamedList + Lower + 'static,
     {
-        self.insert(name, Definition::Func(HostFunc::func_wrap_async(f)))?;
+        self.insert(name, Definition::Func(HostFunc::func_wrap_async(f)?))?;
         Ok(())
     }
 
@@ -550,7 +553,7 @@ impl<T: 'static> LinkerInstance<'_, T> {
         if !self.engine.tunables().concurrency_support {
             bail!("concurrent host functions require `Config::concurrency_support`");
         }
-        self.insert(name, Definition::Func(HostFunc::func_wrap_concurrent(f)))?;
+        self.insert(name, Definition::Func(HostFunc::func_wrap_concurrent(f)?))?;
         Ok(())
     }
 
@@ -661,7 +664,7 @@ impl<T: 'static> LinkerInstance<'_, T> {
         + Sync
         + 'static,
     ) -> Result<()> {
-        self.insert(name, Definition::Func(HostFunc::func_new(func)))?;
+        self.insert(name, Definition::Func(HostFunc::func_new(func)?))?;
         Ok(())
     }
 
@@ -684,7 +687,7 @@ impl<T: 'static> LinkerInstance<'_, T> {
             + Sync
             + 'static,
     {
-        self.insert(name, Definition::Func(HostFunc::func_new_async(func)))?;
+        self.insert(name, Definition::Func(HostFunc::func_new_async(func)?))?;
         Ok(())
     }
 
@@ -712,7 +715,7 @@ impl<T: 'static> LinkerInstance<'_, T> {
         if !self.engine.tunables().concurrency_support {
             bail!("concurrent host functions require `Config::concurrency_support`");
         }
-        self.insert(name, Definition::Func(HostFunc::func_new_concurrent(f)))?;
+        self.insert(name, Definition::Func(HostFunc::func_new_concurrent(f)?))?;
         Ok(())
     }
 
@@ -841,7 +844,7 @@ impl<T: 'static> LinkerInstance<'_, T> {
         Ok(self)
     }
 
-    fn insert(&mut self, name: &str, item: Definition) -> Result<usize> {
+    fn insert(&mut self, name: &str, item: Definition) -> Result<Atom> {
         self.map
             .insert(name, self.strings, self.allow_shadowing, item)
     }
@@ -850,23 +853,3 @@ impl<T: 'static> LinkerInstance<'_, T> {
         self.map.get(name, self.strings)
     }
 }
-
-impl NameMapIntern for Strings {
-    type Key = usize;
-    type BorrowedKey = usize;
-
-    fn intern(&mut self, string: &str) -> Result<usize, OutOfMemory> {
-        if let Some(idx) = self.string2idx.get(string) {
-            return Ok(*idx);
-        }
-        let string: Arc<str> = string.into();
-        let idx = self.strings.len();
-        self.strings.push(string.clone());
-        self.string2idx.insert(string, idx);
-        Ok(idx)
-    }
-
-    fn lookup(&self, string: &str) -> Option<TryCow<'_, usize>> {
-        self.string2idx.get(string).copied().map(TryCow::Owned)
-    }
-}

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

@@ -1,7 +1,7 @@
 use crate::Module;
 use crate::component::ResourceType;
 use crate::component::func::HostFunc;
-use crate::component::linker::{Definition, Strings};
+use crate::component::linker::Definition;
 use crate::component::types::{FutureType, StreamType};
 use crate::runtime::vm::component::ComponentInstance;
 use crate::types::matching;
@@ -13,11 +13,12 @@ use wasmtime_environ::component::{
     TypeStreamTableIndex,
 };
 use wasmtime_environ::prelude::TryPrimaryMap;
+use wasmtime_environ::{Atom, StringPool};
 
 pub struct TypeChecker<'a> {
     pub engine: &'a Engine,
     pub types: &'a Arc<ComponentTypes>,
-    pub strings: &'a Strings,
+    pub strings: &'a StringPool,
     pub imported_resources: Arc<TryPrimaryMap<ResourceIndex, ResourceType>>,
 }
 
@@ -147,7 +148,7 @@ impl TypeChecker<'_> {
     fn instance(
         &mut self,
         expected: &TypeComponentInstance,
-        actual: Option<&NameMap<usize, Definition>>,
+        actual: Option<&NameMap<Atom, Definition>>,
     ) -> Result<()> {
         // Like modules, every export in the expected type must be present in
         // the actual type. It's ok, though, to have extra exports in the actual