Force callers to handle allocation failure in Slab (#12450)

* Add `Vec::drain` method

* Force callers to handle allocation failure in `Slab`

Author: fitzgen

crates/core/src/slab.rs

@@ -14,14 +14,15 @@
 //! # Example
 //!
 //! ```
+//! # fn foo() -> wasmtime_internal_core::error::Result<()> {
 //! use wasmtime_internal_core::slab::Slab;
 //!
 //! let mut slab = Slab::new();
 //!
 //! // Insert some values into the slab.
-//! let rza = slab.alloc("Robert Fitzgerald Diggs");
-//! let gza = slab.alloc("Gary Grice");
-//! let bill = slab.alloc("Bill Gates");
+//! let rza = slab.alloc("Robert Fitzgerald Diggs")?;
+//! let gza = slab.alloc("Gary Grice")?;
+//! let bill = slab.alloc("Bill Gates")?;
 //!
 //! // Allocated elements can be accessed infallibly via indexing (and missing and
 //! // deallocated entries will panic).
@@ -39,7 +40,10 @@
 //! slab.dealloc(bill);
 //!
 //! // Allocate a new entry.
-//! let bill = slab.alloc("Bill Murray");
+//! let bill = slab.alloc("Bill Murray")?;
+//! # Ok(())
+//! # }
+//! # foo().unwrap();
 //! ```
 //!
 //! # Using `Id`s with the Wrong `Slab`
@@ -78,6 +82,8 @@
 //! the following:
 //!
 //! ```rust
+//! use wasmtime_internal_core::error::OutOfMemory;
+//!
 //! pub struct GenerationalId {
 //!     id: wasmtime_internal_core::slab::Id,
 //!     generation: u32,
@@ -94,10 +100,10 @@
 //! }
 //!
 //! impl<T> GenerationalSlab<T> {
-//!     pub fn alloc(&mut self, value: T) -> GenerationalId {
+//!     pub fn alloc(&mut self, value: T) -> Result<GenerationalId, OutOfMemory> {
 //!         let generation = self.generation;
-//!         let id = self.slab.alloc(GenerationalEntry { value, generation });
-//!         GenerationalId { id, generation }
+//!         let id = self.slab.alloc(GenerationalEntry { value, generation })?;
+//!         Ok(GenerationalId { id, generation })
 //!     }
 //!
 //!     pub fn get(&self, id: GenerationalId) -> Option<&T> {
@@ -127,9 +133,10 @@
 //! }
 //! ```
 
+use crate::alloc::Vec;
+use crate::error::OutOfMemory;
 use core::fmt;
 use core::num::NonZeroU32;
-use std_alloc::vec::Vec;
 
 /// An identifier for an allocated value inside a `slab`.
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -255,10 +262,10 @@ impl<T> Slab<T> {
     /// Construct a new, empty slab, pre-reserving space for at least `capacity`
     /// elements.
     #[inline]
-    pub fn with_capacity(capacity: usize) -> Self {
+    pub fn with_capacity(capacity: usize) -> Result<Self, OutOfMemory> {
         let mut slab = Self::new();
-        slab.reserve(capacity);
-        slab
+        slab.reserve(capacity)?;
+        Ok(slab)
     }
 
     /// Ensure that there is space for at least `additional` elements in this
@@ -267,29 +274,31 @@ impl<T> Slab<T> {
     /// # Panics
     ///
     /// Panics if the new capacity exceeds `Self::MAX_CAPACITY`.
-    pub fn reserve(&mut self, additional: usize) {
+    pub fn reserve(&mut self, additional: usize) -> Result<(), OutOfMemory> {
         let cap = self.capacity();
         let len = self.len();
         assert!(cap >= len);
         if cap - len >= additional {
             // Already have `additional` capacity available.
-            return;
+            return Ok(());
         }
 
-        self.entries.reserve(additional);
+        self.entries.reserve(additional)?;
 
         // Maintain the invariant that `i <= MAX_CAPACITY` for all indices `i`
         // in `self.entries`.
         assert!(self.entries.capacity() <= Self::MAX_CAPACITY);
+
+        Ok(())
     }
 
-    fn double_capacity(&mut self) {
+    fn double_capacity(&mut self) -> Result<(), OutOfMemory> {
         // Double our capacity to amortize the cost of resizing. But make sure
         // we add some amount of minimum additional capacity, since doubling
         // zero capacity isn't useful.
         const MIN_CAPACITY: usize = 16;
         let additional = core::cmp::max(self.entries.capacity(), MIN_CAPACITY);
-        self.reserve(additional);
+        self.reserve(additional)
     }
 
     /// What is the capacity of this slab? That is, how many entries can it
@@ -336,7 +345,9 @@ impl<T> Slab<T> {
         self.free.take().or_else(|| {
             if self.entries.len() < self.entries.capacity() {
                 let index = EntryIndex::new(self.entries.len());
-                self.entries.push(Entry::Free { next_free: None });
+                self.entries
+                    .push(Entry::Free { next_free: None })
+                    .expect("have capacity");
                 Some(index)
             } else {
                 None
@@ -352,9 +363,9 @@ impl<T> Slab<T> {
     /// Panics if allocating this value requires reallocating the underlying
     /// storage, and the new capacity exceeds `Slab::MAX_CAPACITY`.
     #[inline]
-    pub fn alloc(&mut self, value: T) -> Id {
+    pub fn alloc(&mut self, value: T) -> Result<Id, OutOfMemory> {
         self.try_alloc(value)
-            .unwrap_or_else(|value| self.alloc_slow(value))
+            .or_else(|value| self.alloc_slow(value))
     }
 
     /// Get the `Id` that will be returned for the next allocation in this slab.
@@ -366,12 +377,13 @@ impl<T> Slab<T> {
 
     #[inline(never)]
     #[cold]
-    fn alloc_slow(&mut self, value: T) -> Id {
+    fn alloc_slow(&mut self, value: T) -> Result<Id, OutOfMemory> {
         // Reserve additional capacity, since we didn't have space for the
         // allocation.
-        self.double_capacity();
+        self.double_capacity()?;
         // After which the allocation will succeed.
-        self.try_alloc(value).ok().unwrap()
+        let id = self.try_alloc(value).ok().unwrap();
+        Ok(id)
     }
 
     /// Get a shared borrow of the value associated with `id`.

crates/wasmtime/src/runtime/gc/enabled/rooting.rs

@@ -164,6 +164,7 @@
 use crate::runtime::vm::{GcRootsList, GcStore, VMGcRef};
 use crate::{
     AsContext, AsContextMut, GcRef, Result, RootedGcRef,
+    error::OutOfMemory,
     store::{AsStoreOpaque, AutoAssertNoGc, StoreId, StoreOpaque},
 };
 use crate::{ValRaw, prelude::*};
@@ -1047,7 +1048,7 @@ impl<T: GcRef> Rooted<T> {
     pub(crate) fn _to_owned_rooted(&self, store: &mut StoreOpaque) -> Result<OwnedRooted<T>> {
         let mut store = AutoAssertNoGc::new(store);
         let gc_ref = self.try_clone_gc_ref(&mut store)?;
-        Ok(OwnedRooted::new(&mut store, gc_ref))
+        Ok(OwnedRooted::new(&mut store, gc_ref)?)
     }
 
     /// Are these two `Rooted<T>`s the same GC root?
@@ -1627,7 +1628,10 @@ where
     /// `gc_ref` should be a GC reference pointing to an instance of the GC type
     /// that `T` represents. Failure to uphold this invariant is memory safe but
     /// will result in general incorrectness such as panics and wrong results.
-    pub(crate) fn new(store: &mut AutoAssertNoGc<'_>, gc_ref: VMGcRef) -> Self {
+    pub(crate) fn new(
+        store: &mut AutoAssertNoGc<'_>,
+        gc_ref: VMGcRef,
+    ) -> Result<Self, OutOfMemory> {
         // We always have the opportunity to trim and unregister stale
         // owned roots whenever we have a mut borrow to the store. We
         // take the opportunity to do so here to avoid tying growth of
@@ -1639,20 +1643,20 @@ where
         store.trim_gc_liveness_flags(false);
 
         let roots = store.gc_roots_mut();
-        let id = roots.owned_rooted.alloc(gc_ref);
+        let id = roots.owned_rooted.alloc(gc_ref)?;
         let liveness_flag = Arc::new(());
         roots
             .liveness_flags
             .push((Arc::downgrade(&liveness_flag), id));
-        OwnedRooted {
+        Ok(OwnedRooted {
             inner: GcRootIndex {
                 store_id: store.id(),
                 generation: 0,
                 index: PackedIndex::new_owned(id),
             },
             liveness_flag,
             _phantom: marker::PhantomData,
-        }
+        })
     }
 
     #[inline]

crates/wasmtime/src/runtime/store.rs

@@ -2780,13 +2780,18 @@ at https://bytecodealliance.org/security.
     /// Get an owned rooted reference to the pending exception,
     /// without taking it off the store.
     #[cfg(all(feature = "gc", feature = "debug"))]
-    pub(crate) fn pending_exception_owned_rooted(&mut self) -> Option<OwnedRooted<ExnRef>> {
+    pub(crate) fn pending_exception_owned_rooted(
+        &mut self,
+    ) -> Result<Option<OwnedRooted<ExnRef>>, crate::error::OutOfMemory> {
         let mut nogc = AutoAssertNoGc::new(self);
-        nogc.pending_exception.take().map(|vmexnref| {
-            let cloned = nogc.clone_gc_ref(vmexnref.as_gc_ref());
-            nogc.pending_exception = Some(cloned.into_exnref_unchecked());
-            OwnedRooted::new(&mut nogc, vmexnref.into())
-        })
+        nogc.pending_exception
+            .take()
+            .map(|vmexnref| {
+                let cloned = nogc.clone_gc_ref(vmexnref.as_gc_ref());
+                nogc.pending_exception = Some(cloned.into_exnref_unchecked());
+                OwnedRooted::new(&mut nogc, vmexnref.into())
+            })
+            .transpose()
     }
 
     #[cfg(feature = "gc")]

crates/wasmtime/src/runtime/type_registry.rs

@@ -738,6 +738,13 @@ impl TypeRegistryInner {
 
         log::trace!("hash-consing map miss: making new registration");
 
+        // Reserve capacity before making any side effects (like incrementing
+        // other rec groups ref counts). We do not want to panic or
+        // `?`-propagate an error when we have applied partial side effects.
+        //
+        // TODO(#12069): handle allocation failure here.
+        self.types.reserve(non_canon_types.len()).unwrap();
+
         // Inter-group edges: increment the referenced group's ref
         // count, because these other rec groups shouldn't be dropped
         // while this rec group is still alive.
@@ -760,7 +767,9 @@ impl TypeRegistryInner {
         let shared_type_indices: Box<[_]> = non_canon_types
             .iter()
             .map(|(module_index, ty)| {
-                let engine_index = slab_id_to_shared_type_index(self.types.alloc(None));
+                let engine_index = slab_id_to_shared_type_index(
+                    self.types.alloc(None).expect("reserved capacity"),
+                );
                 log::trace!(
                     "reserved {engine_index:?} for {module_index:?} = non-canonical {ty:?}"
                 );

crates/wasmtime/src/runtime/vm/gc/func_ref.rs

@@ -50,10 +50,10 @@ impl FuncRefTable {
     /// The given `func_ref` must point to a valid `VMFuncRef` and must remain
     /// valid for the duration of this table's lifetime.
     pub unsafe fn intern(&mut self, func_ref: Option<SendSyncPtr<VMFuncRef>>) -> FuncRefTableId {
-        *self
-            .interned
-            .entry(func_ref)
-            .or_insert_with(|| FuncRefTableId(self.slab.alloc(func_ref)))
+        *self.interned.entry(func_ref).or_insert_with(|| {
+            // TODO(#12069): handle allocation failure here
+            FuncRefTableId(self.slab.alloc(func_ref).unwrap())
+        })
     }
 
     /// Get the `VMFuncRef` associated with the given ID.

crates/wasmtime/src/runtime/vm/gc/host_data.rs

@@ -39,7 +39,8 @@ fn deref_box_mut<T: ?Sized>(b: &mut Box<T>) -> &mut T {
 impl ExternRefHostDataTable {
     /// Allocate a new `externref` host data value.
     pub fn alloc(&mut self, value: Box<dyn Any + Send + Sync>) -> ExternRefHostDataId {
-        let id = self.slab.alloc(value);
+        // TODO(#12069): handle allocation failure here
+        let id = self.slab.alloc(value).unwrap();
         let id = ExternRefHostDataId(id);
         log::trace!("allocated new externref host data: {id:?}");
         id

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

@@ -837,6 +837,8 @@ impl CallThreadState {
                     let exn = store
                         .as_store_opaque()
                         .pending_exception_owned_rooted()
+                        // TODO(#12069): handle allocation failure here
+                        .unwrap()
                         .expect("exception should be set when we are throwing");
                     store.block_on_debug_handler(crate::DebugEvent::CaughtExceptionThrown(exn))
                 }
@@ -848,6 +850,8 @@ impl CallThreadState {
                     let exn = store
                         .as_store_opaque()
                         .pending_exception_owned_rooted()
+                        // TODO(#12069): handle allocation failure here
+                        .unwrap()
                         .expect("exception should be set when we are throwing");
                     store.block_on_debug_handler(crate::DebugEvent::UncaughtExceptionThrown(
                         exn.clone(),