Introduce wasmtime::Store::try_new, which handles OOM (#12415)

* Introduce `wasmtime::Store::try_new`, which handles OOM

`Store::new` is an infallible constructor, so there is not a direct way to make
it return an error on OOM. Additionally, it is one of the most-used functions in
the Wasmtime embedder API, so changing its signature to return a `Result` is a
non-starter -- it would cause way too much pain. So instead we define
`Store::try_new` which returns a `Result` and make `Store::new` call and unwrap
that new constructor.

Part of #12069

* update disas tests and fix winch

* Disable concurrency support in `Store::try_new` OOM test

* Add attributes that were lost in rebase

Author: fitzgen

crates/environ/src/builtin.rs

@@ -11,7 +11,7 @@ macro_rules! foreach_builtin_function {
             // Returns an index for wasm's `table.init`.
             table_init(vmctx: vmctx, table: u32, elem: u32, dst: u64, src: u64, len: u64) -> bool;
             // Returns an index for wasm's `elem.drop`.
-            elem_drop(vmctx: vmctx, elem: u32);
+            elem_drop(vmctx: vmctx, elem: u32) -> bool;
             // Returns an index for wasm's `memory.copy`
             memory_copy(vmctx: vmctx, dst_index: u32, dst: u64, src_index: u32, src: u64, len: u64) -> bool;
             // Returns an index for wasm's `memory.fill` instruction.
@@ -21,7 +21,7 @@ macro_rules! foreach_builtin_function {
             // Returns a value for wasm's `ref.func` instruction.
             ref_func(vmctx: vmctx, func: u32) -> pointer;
             // Returns an index for wasm's `data.drop` instruction.
-            data_drop(vmctx: vmctx, data: u32);
+            data_drop(vmctx: vmctx, data: u32) -> bool;
             // Returns a table entry after lazily initializing it.
             table_get_lazy_init_func_ref(vmctx: vmctx, table: u32, index: u64) -> pointer;
             // Returns an index for Wasm's `table.grow` instruction for `funcref`s.

crates/fuzzing/tests/oom.rs

@@ -138,6 +138,19 @@ fn linker_new() -> Result<()> {
     })
 }
 
+#[test]
+#[cfg(arc_try_new)]
+fn store_try_new() -> Result<()> {
+    let mut config = Config::new();
+    config.enable_compiler(false);
+    config.concurrency_support(false);
+    let engine = Engine::new(&config)?;
+    OomTest::new().test(|| {
+        let _ = Store::try_new(&engine, ())?;
+        Ok(())
+    })
+}
+
 fn ok_if_not_oom(error: Error) -> Result<()> {
     if error.is::<OutOfMemory>() {
         Err(error)

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

@@ -758,7 +758,7 @@ impl<'a> Instantiator<'a> {
             Arc::new(imported_resources),
             imports,
             store.traitobj(),
-        );
+        )?;
         let id = store.store_data_mut().push_component_instance(instance);
 
         Ok(Instantiator {

crates/wasmtime/src/runtime/store.rs

@@ -85,6 +85,7 @@ use crate::ThrownException;
 use crate::component::ComponentStoreData;
 #[cfg(feature = "component-model")]
 use crate::component::concurrent;
+use crate::error::OutOfMemory;
 #[cfg(feature = "async")]
 use crate::fiber;
 use crate::module::RegisteredModuleId;
@@ -474,7 +475,7 @@ pub struct StoreOpaque {
     #[cfg(feature = "stack-switching")]
     continuations: Vec<Box<VMContRef>>,
 
-    instances: PrimaryMap<InstanceId, StoreInstance>,
+    instances: wasmtime_environ::collections::PrimaryMap<InstanceId, StoreInstance>,
 
     #[cfg(feature = "component-model")]
     num_component_instances: usize,
@@ -725,6 +726,13 @@ impl<T> Store<T> {
     /// tables created to 10,000. This can be overridden with the
     /// [`Store::limiter`] configuration method.
     pub fn new(engine: &Engine, data: T) -> Self {
+        Self::try_new(engine, data).expect(
+            "allocation failure during `Store::new` (use `Store::try_new` to handle such errors)",
+        )
+    }
+
+    /// Like `Store::new` but returns an error on allocation failure.
+    pub fn try_new(engine: &Engine, data: T) -> Result<Self> {
         let store_data = StoreData::new();
         log::trace!("creating new store {:?}", store_data.id());
 
@@ -736,7 +744,7 @@ impl<T> Store<T> {
             vm_store_context: Default::default(),
             #[cfg(feature = "stack-switching")]
             continuations: Vec::new(),
-            instances: PrimaryMap::new(),
+            instances: wasmtime_environ::collections::PrimaryMap::new(),
             #[cfg(feature = "component-model")]
             num_component_instances: 0,
             signal_handler: None,
@@ -790,7 +798,7 @@ impl<T> Store<T> {
             #[cfg(feature = "debug")]
             breakpoints: Default::default(),
         };
-        let mut inner = Box::new(StoreInner {
+        let mut inner = try_new::<Box<_>>(StoreInner {
             inner,
             limiter: None,
             call_hook: None,
@@ -799,7 +807,7 @@ impl<T> Store<T> {
             data_no_provenance: ManuallyDrop::new(data),
             #[cfg(feature = "debug")]
             debug_handler: None,
-        });
+        })?;
 
         let store_data =
             <NonNull<ManuallyDrop<T>>>::from(&mut inner.data_no_provenance).cast::<()>();
@@ -825,22 +833,30 @@ impl<T> Store<T> {
             // (also no limiter is passed in) so it won't have an async await
             // point meaning that it should be ok to assert the future is
             // always ready.
-            let id = vm::assert_ready(inner.allocate_instance(
+            let result = vm::assert_ready(inner.allocate_instance(
                 None,
                 AllocateInstanceKind::Dummy {
                     allocator: &allocator,
                 },
                 info,
                 Default::default(),
-            ))
-            .expect("failed to allocate default callee");
+            ));
+            let id = match result {
+                Ok(id) => id,
+                Err(e) => {
+                    if e.is::<OutOfMemory>() {
+                        return Err(e);
+                    }
+                    panic!("instance allocator failed to allocate default callee")
+                }
+            };
             let default_caller_vmctx = inner.instance(id).vmctx();
             inner.default_caller_vmctx = default_caller_vmctx.into();
         }
 
-        Self {
+        Ok(Self {
             inner: ManuallyDrop::new(inner),
-        }
+        })
     }
 
     /// Access the underlying `T` data owned by this `Store`.
@@ -2712,7 +2728,7 @@ at https://bytecodealliance.org/security.
                 self.instances.push(StoreInstance {
                     handle,
                     kind: StoreInstanceKind::Real { module_id },
-                })
+                })?
             }
             AllocateInstanceKind::Dummy { .. } => {
                 log::trace!(
@@ -2722,7 +2738,7 @@ at https://bytecodealliance.org/security.
                 self.instances.push(StoreInstance {
                     handle,
                     kind: StoreInstanceKind::Dummy,
-                })
+                })?
             }
         };
 

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

@@ -28,6 +28,7 @@ use core::mem::offset_of;
 use core::pin::Pin;
 use core::ptr::NonNull;
 use wasmtime_environ::component::*;
+use wasmtime_environ::error::OutOfMemory;
 use wasmtime_environ::{HostPtr, PrimaryMap, VMSharedTypeIndex};
 
 #[allow(
@@ -317,7 +318,7 @@ impl ComponentInstance {
         resource_types: Arc<PrimaryMap<ResourceIndex, ResourceType>>,
         imports: &Arc<PrimaryMap<RuntimeImportIndex, RuntimeImport>>,
         store: NonNull<dyn VMStore>,
-    ) -> OwnedComponentInstance {
+    ) -> 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());
@@ -342,11 +343,11 @@ impl ComponentInstance {
             store: VMStoreRawPtr(store),
             post_return_arg: None,
             vmctx: OwnedVMContext::new(),
-        });
+        })?;
         unsafe {
             ret.get_mut().initialize_vmctx();
         }
-        ret
+        Ok(ret)
     }
 
     #[inline]

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

@@ -34,11 +34,12 @@ use core::sync::atomic::AtomicU64;
 use core::{mem, ptr};
 #[cfg(feature = "gc")]
 use wasmtime_environ::ModuleInternedTypeIndex;
+use wasmtime_environ::error::OutOfMemory;
 use wasmtime_environ::{
     DataIndex, DefinedGlobalIndex, DefinedMemoryIndex, DefinedTableIndex, DefinedTagIndex,
-    ElemIndex, EntityIndex, EntityRef, EntitySet, FuncIndex, GlobalIndex, HostPtr, MemoryIndex,
-    PrimaryMap, PtrSize, TableIndex, TableInitialValue, TableSegmentElements, TagIndex, Trap,
-    VMCONTEXT_MAGIC, VMOffsets, VMSharedTypeIndex, packed_option::ReservedValue,
+    ElemIndex, EntityIndex, EntityRef, FuncIndex, GlobalIndex, HostPtr, MemoryIndex, PrimaryMap,
+    PtrSize, TableIndex, TableInitialValue, TableSegmentElements, TagIndex, Trap, VMCONTEXT_MAGIC,
+    VMOffsets, VMSharedTypeIndex, packed_option::ReservedValue,
 };
 #[cfg(feature = "wmemcheck")]
 use wasmtime_wmemcheck::Wmemcheck;
@@ -167,11 +168,11 @@ impl Instance {
         req: InstanceAllocationRequest,
         memories: PrimaryMap<DefinedMemoryIndex, (MemoryAllocationIndex, Memory)>,
         tables: PrimaryMap<DefinedTableIndex, (TableAllocationIndex, Table)>,
-    ) -> InstanceHandle {
+    ) -> Result<InstanceHandle, OutOfMemory> {
         let module = req.runtime_info.env_module();
         let memory_tys = &module.memories;
-        let dropped_elements = EntitySet::with_capacity(module.passive_elements.len());
-        let dropped_data = EntitySet::with_capacity(module.passive_data_map.len());
+        let dropped_elements = EntitySet::with_capacity(module.passive_elements.len())?;
+        let dropped_data = EntitySet::with_capacity(module.passive_data_map.len())?;
 
         #[cfg(feature = "wmemcheck")]
         let wmemcheck_state = if req.store.engine().config().wmemcheck {
@@ -200,14 +201,14 @@ impl Instance {
             wmemcheck_state,
             store: None,
             vmctx: OwnedVMContext::new(),
-        });
+        })?;
 
         // SAFETY: this vmctx was allocated with the same layout above, so it
         // should be safe to initialize with the same values here.
         unsafe {
             ret.get_mut().initialize_vmctx(req.store, req.imports);
         }
-        ret
+        Ok(ret)
     }
 
     /// Converts a raw `VMContext` pointer into a raw `Instance` pointer.
@@ -1050,13 +1051,18 @@ impl Instance {
     }
 
     /// Drop an element.
-    pub(crate) fn elem_drop(self: Pin<&mut Self>, elem_index: ElemIndex) {
+    pub(crate) fn elem_drop(
+        self: Pin<&mut Self>,
+        elem_index: ElemIndex,
+    ) -> Result<(), OutOfMemory> {
         // https://webassembly.github.io/reference-types/core/exec/instructions.html#exec-elem-drop
 
-        self.dropped_elements_mut().insert(elem_index);
+        self.dropped_elements_mut().insert(elem_index)?;
 
         // Note that we don't check that we actually removed a segment because
         // dropping a non-passive segment is a no-op (not a trap).
+
+        Ok(())
     }
 
     /// Get a locally-defined memory.
@@ -1218,11 +1224,16 @@ impl Instance {
     }
 
     /// Drop the given data segment, truncating its length to zero.
-    pub(crate) fn data_drop(self: Pin<&mut Self>, data_index: DataIndex) {
-        self.dropped_data_mut().insert(data_index);
+    pub(crate) fn data_drop(
+        self: Pin<&mut Self>,
+        data_index: DataIndex,
+    ) -> Result<(), OutOfMemory> {
+        self.dropped_data_mut().insert(data_index)?;
 
         // Note that we don't check that we actually removed a segment because
         // dropping a non-passive segment is a no-op (not a trap).
+
+        Ok(())
     }
 
     /// Get a table by index regardless of whether it is locally-defined
@@ -1891,7 +1902,7 @@ impl<T: InstanceLayout> OwnedInstance<T> {
     /// Allocates a new `OwnedInstance` and places `instance` inside of it.
     ///
     /// This will `instance`
-    pub(super) fn new(mut instance: T) -> OwnedInstance<T> {
+    pub(super) fn new(mut instance: T) -> Result<OwnedInstance<T>, OutOfMemory> {
         let layout = instance.layout();
         debug_assert!(layout.size() >= size_of_val(&instance));
         debug_assert!(layout.align() >= align_of_val(&instance));
@@ -1906,10 +1917,9 @@ impl<T: InstanceLayout> OwnedInstance<T> {
                 alloc::alloc::alloc(layout)
             }
         };
-        if ptr.is_null() {
-            alloc::alloc::handle_alloc_error(layout);
-        }
-        let instance_ptr = NonNull::new(ptr.cast::<T>()).unwrap();
+        let Some(instance_ptr) = NonNull::new(ptr.cast::<T>()) else {
+            return Err(OutOfMemory::new(layout.size()));
+        };
 
         // SAFETY: it's part of the unsafe contract of `InstanceLayout` that the
         // `add` here is appropriate for the layout allocated.
@@ -1938,7 +1948,7 @@ impl<T: InstanceLayout> OwnedInstance<T> {
         );
         debug_assert_eq!(vmctx_self_reference.addr(), ret.get().vmctx().addr());
 
-        ret
+        Ok(ret)
     }
 
     /// Gets the raw underlying `&Instance` from this handle.

crates/wasmtime/src/runtime/vm/instance/allocator.rs

@@ -344,7 +344,7 @@ impl dyn InstanceAllocator + '_ {
                 request,
                 mem::take(&mut guard.memories),
                 mem::take(&mut guard.tables),
-            ))
+            )?)
         };
 
         struct DeallocateOnDrop<'a> {

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

@@ -541,9 +541,10 @@ fn table_init(
 }
 
 // Implementation of `elem.drop`.
-fn elem_drop(store: &mut dyn VMStore, instance: InstanceId, elem_index: u32) {
+fn elem_drop(store: &mut dyn VMStore, instance: InstanceId, elem_index: u32) -> Result<()> {
     let elem_index = ElemIndex::from_u32(elem_index);
-    store.instance_mut(instance).elem_drop(elem_index)
+    store.instance_mut(instance).elem_drop(elem_index)?;
+    Ok(())
 }
 
 // Implementation of `memory.copy`.
@@ -606,9 +607,10 @@ fn ref_func(store: &mut dyn VMStore, instance: InstanceId, func_index: u32) -> N
 }
 
 // Implementation of `data.drop`.
-fn data_drop(store: &mut dyn VMStore, instance: InstanceId, data_index: u32) {
+fn data_drop(store: &mut dyn VMStore, instance: InstanceId, data_index: u32) -> Result<()> {
     let data_index = DataIndex::from_u32(data_index);
-    store.instance_mut(instance).data_drop(data_index)
+    store.instance_mut(instance).data_drop(data_index)?;
+    Ok(())
 }
 
 // Returns a table entry after lazily initializing it.

tests/disas/passive-data.wat

@@ -40,13 +40,13 @@
 ;;     gv1 = load.i64 notrap aligned readonly gv0+8
 ;;     gv2 = load.i64 notrap aligned gv1+16
 ;;     gv3 = vmctx
-;;     sig0 = (i64 vmctx, i32) tail
+;;     sig0 = (i64 vmctx, i32) -> i8 tail
 ;;     fn0 = colocated u805306368:8 sig0
 ;;     stack_limit = gv2
 ;;
 ;;                                 block0(v0: i64, v1: i64):
 ;; @0044                               v2 = iconst.i32 0
-;; @0044                               call fn0(v0, v2)  ; v2 = 0
+;; @0044                               v4 = call fn0(v0, v2)  ; v2 = 0
 ;; @0047                               jump block1
 ;;
 ;;                                 block1:

tests/disas/winch/x64/table/init_copy_drop.wat

@@ -243,7 +243,7 @@
 ;;       movq    %r14, %rdi
 ;;       movl    $0, %esi
 ;;       movl    0xc(%rsp), %edx
-;;       callq   0x9c0
+;;       callq   0x9df
 ;;       addq    $0xc, %rsp
 ;;       addq    $4, %rsp
 ;;       movq    0x18(%rsp), %r14