Use TryVec in the store's FuncRefs structure

Author: fitzgen

crates/wasmtime/src/runtime/func.rs

@@ -381,7 +381,7 @@ impl Func {
 
         // SAFETY: the `T` used by `func` matches the `T` of the store we're
         // inserting into via this function's type signature.
-        unsafe { host.into_func(store) }
+        unsafe { host.into_func(store).panic_on_oom() }
     }
 
     /// Creates a new [`Func`] with the given arguments, although has fewer
@@ -425,7 +425,7 @@ impl Func {
 
         // SAFETY: the `T` used by `func` matches the `T` of the store we're
         // inserting into via this function's type signature.
-        unsafe { host.into_func(store) }
+        unsafe { host.into_func(store).panic_on_oom() }
     }
 
     /// Creates a new host-defined WebAssembly function which, when called,
@@ -516,7 +516,7 @@ impl Func {
 
         // SAFETY: the `T` used by `func` matches the `T` of the store we're
         // inserting into via this function's type signature.
-        unsafe { host.into_func(store) }
+        unsafe { host.into_func(store).panic_on_oom() }
     }
 
     /// Creates a new `Func` from a store and a funcref within that store.
@@ -795,7 +795,7 @@ impl Func {
 
         // SAFETY: The `T` the closure takes is the same as the `T` of the store
         // we're inserting into via the type signature above.
-        unsafe { host.into_func(store) }
+        unsafe { host.into_func(store).panic_on_oom() }
     }
 
     /// Same as [`Func::wrap`], except the closure asynchronously produces the
@@ -817,7 +817,7 @@ impl Func {
 
         // SAFETY: The `T` the closure takes is the same as the `T` of the store
         // we're inserting into via the type signature above.
-        unsafe { host.into_func(store) }
+        unsafe { host.into_func(store).panic_on_oom() }
     }
 
     /// Returns the underlying wasm type that this `Func` has.
@@ -2625,13 +2625,13 @@ impl HostFunc {
     ///
     /// Can only be inserted into stores with a matching `T` relative to when
     /// this `HostFunc` was first created.
-    pub unsafe fn to_func(self: &Arc<Self>, store: &mut StoreOpaque) -> Func {
+    pub unsafe fn to_func(self: &Arc<Self>, store: &mut StoreOpaque) -> Result<Func, OutOfMemory> {
         self.validate_store(store);
         let (funcrefs, modules) = store.func_refs_and_modules();
-        let funcref = funcrefs.push_arc_host(self.clone(), modules);
+        let funcref = funcrefs.push_arc_host(self.clone(), modules)?;
         // SAFETY: this funcref was just pushed within the store, so it's safe
         // to say this store owns it.
-        unsafe { Func::from_vm_func_ref(store.id(), funcref) }
+        Ok(unsafe { Func::from_vm_func_ref(store.id(), funcref) })
     }
 
     /// Inserts this `HostFunc` into a `Store`, returning the `Func` pointing to
@@ -2685,18 +2685,18 @@ impl HostFunc {
     }
 
     /// Same as [`HostFunc::to_func`], different ownership.
-    unsafe fn into_func(self, store: &mut StoreOpaque) -> Func {
+    unsafe fn into_func(self, store: &mut StoreOpaque) -> Result<Func, OutOfMemory> {
         self.validate_store(store);
 
         // This function could be called by a guest at any time, and it requires
         // fibers, so the store now required async entrypoints.
         store.set_async_required(self.asyncness);
 
         let (funcrefs, modules) = store.func_refs_and_modules();
-        let funcref = funcrefs.push_box_host(Box::new(self), modules);
+        let funcref = funcrefs.push_box_host(Box::new(self), modules)?;
         // SAFETY: this funcref was just pushed within `store`, so it's safe to
         // say it's owned by the store's id.
-        unsafe { Func::from_vm_func_ref(store.id(), funcref) }
+        Ok(unsafe { Func::from_vm_func_ref(store.id(), funcref) })
     }
 
     fn validate_store(&self, store: &mut StoreOpaque) {

crates/wasmtime/src/runtime/instance.rs

@@ -976,14 +976,14 @@ fn pre_instantiate_raw(
         // will insert a function into the store automatically as part of
         // instantiation, so reserve space here to make insertion more efficient
         // as it won't have to realloc during the instantiation.
-        funcrefs.reserve_storage(host_funcs);
+        funcrefs.reserve_storage(host_funcs)?;
 
         // The usage of `to_extern_store_rooted` requires that the items are
         // rooted via another means, which happens here by cloning the list of
         // items into the store once. This avoids cloning each individual item
         // below.
-        funcrefs.push_instance_pre_definitions(items.clone());
-        funcrefs.push_instance_pre_func_refs(func_refs.clone());
+        funcrefs.push_instance_pre_definitions(items.clone())?;
+        funcrefs.push_instance_pre_func_refs(func_refs.clone())?;
     }
 
     store.set_async_required(asyncness);

crates/wasmtime/src/runtime/linker.rs

@@ -13,7 +13,7 @@ use core::future::Future;
 use core::marker;
 use core::mem::MaybeUninit;
 use log::warn;
-use wasmtime_environ::{Atom, StringPool};
+use wasmtime_environ::{Atom, PanicOnOom, StringPool};
 
 /// Structure used to link wasm modules/instances together.
 ///
@@ -1348,7 +1348,7 @@ impl Definition {
             // SAFETY: the contract of this function is the same as what's
             // required of `to_func`, that `T` of the store matches the `T` of
             // this original definition.
-            Definition::HostFunc(func) => unsafe { func.to_func(store).into() },
+            Definition::HostFunc(func) => unsafe { func.to_func(store).panic_on_oom().into() },
         }
     }
 

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

@@ -22,11 +22,11 @@ pub struct FuncRefs {
 
     /// Pointers into `self.bump` for entries that need `wasm_call` field filled
     /// in.
-    with_holes: Vec<SendSyncPtr<VMFuncRef>>,
+    with_holes: TryVec<SendSyncPtr<VMFuncRef>>,
 
     /// General-purpose storage of "function things" that need to live as long
     /// as the entire store.
-    storage: Vec<Storage>,
+    storage: TryVec<Storage>,
 }
 
 /// Various items to place in `FuncRefs::storage`
@@ -90,17 +90,17 @@ impl FuncRefs {
         &mut self,
         func_ref: VMFuncRef,
         modules: &ModuleRegistry,
-    ) -> NonNull<VMFuncRef> {
+    ) -> Result<NonNull<VMFuncRef>, OutOfMemory> {
         debug_assert!(func_ref.wasm_call.is_none());
         let func_ref = self.bump.get_mut().alloc(func_ref);
         // SAFETY: it's a contract of this function itself that `func_ref` has a
         // valid vmctx field to read.
         let has_hole = unsafe { !try_fill(func_ref, modules) };
         let unpatched = SendSyncPtr::from(func_ref);
         if has_hole {
-            self.with_holes.push(unpatched);
+            self.with_holes.push(unpatched)?;
         }
-        unpatched.as_non_null()
+        Ok(unpatched.as_non_null())
     }
 
     /// Patch any `VMFuncRef::wasm_call`s that need filling in.
@@ -110,25 +110,31 @@ impl FuncRefs {
     }
 
     /// Reserves `amt` space for extra items in "storage" for this store.
-    pub fn reserve_storage(&mut self, amt: usize) {
-        self.storage.reserve(amt);
+    pub fn reserve_storage(&mut self, amt: usize) -> Result<(), OutOfMemory> {
+        self.storage.reserve(amt)
     }
 
     /// Push pre-patched `VMFuncRef`s from an `InstancePre`.
     ///
     /// This is used to ensure that the store itself persists the entire list of
     /// `funcs` for the entire lifetime of the store.
-    pub fn push_instance_pre_func_refs(&mut self, funcs: Arc<TryVec<VMFuncRef>>) {
-        self.storage.push(Storage::InstancePreFuncRefs { funcs });
+    pub fn push_instance_pre_func_refs(
+        &mut self,
+        funcs: Arc<TryVec<VMFuncRef>>,
+    ) -> Result<(), OutOfMemory> {
+        self.storage.push(Storage::InstancePreFuncRefs { funcs })
     }
 
     /// Push linker definitions into storage, keeping them alive for the entire
     /// lifetime of the store.
     ///
     /// This is used to keep linker-defined functions' vmctx values alive, for
     /// example.
-    pub fn push_instance_pre_definitions(&mut self, defs: Arc<TryVec<Definition>>) {
-        self.storage.push(Storage::InstancePreDefinitions { defs });
+    pub fn push_instance_pre_definitions(
+        &mut self,
+        defs: Arc<TryVec<Definition>>,
+    ) -> Result<(), OutOfMemory> {
+        self.storage.push(Storage::InstancePreDefinitions { defs })
     }
 
     /// Pushes a shared host function into this store.
@@ -147,25 +153,25 @@ impl FuncRefs {
         &mut self,
         func: Arc<HostFunc>,
         modules: &ModuleRegistry,
-    ) -> NonNull<VMFuncRef> {
+    ) -> Result<NonNull<VMFuncRef>, OutOfMemory> {
         debug_assert!(func.func_ref().wasm_call.is_none());
         // SAFETY: the vmctx field in the funcref of `HostFunc` is safe to read.
-        let ret = unsafe { self.push(func.func_ref().clone(), modules) };
-        self.storage.push(Storage::ArcHost { func });
-        ret
+        let ret = unsafe { self.push(func.func_ref().clone(), modules)? };
+        self.storage.push(Storage::ArcHost { func })?;
+        Ok(ret)
     }
 
     /// Same as `push_arc_host`, but for owned host functions.
     pub fn push_box_host(
         &mut self,
         func: Box<HostFunc>,
         modules: &ModuleRegistry,
-    ) -> NonNull<VMFuncRef> {
+    ) -> Result<NonNull<VMFuncRef>, OutOfMemory> {
         debug_assert!(func.func_ref().wasm_call.is_none());
         // SAFETY: the vmctx field in the funcref of `HostFunc` is safe to read.
-        let ret = unsafe { self.push(func.func_ref().clone(), modules) };
-        self.storage.push(Storage::BoxHost { func });
-        ret
+        let ret = unsafe { self.push(func.func_ref().clone(), modules)? };
+        self.storage.push(Storage::BoxHost { func })?;
+        Ok(ret)
     }
 }