Fix corrupted pooling allocator state on OOM (#13579)

This commit fixes an issue where if the pooling allocator generated OOM
at just the right time then it would leave the pooling allocator in a
corrupt state where the instance count was incremented, never
decremented. This additionally fixes a separate issue where if an OOM
happened just a bit later then it would mean that tables/memories from
the pooling allocator would leak and not be properly destroyed within
the allocator, leaking their slots.

Author: alexcrichton

crates/fuzzing/tests/oom/instance.rs

@@ -1,8 +1,8 @@
 #![cfg(arc_try_new)]
 
 use wasmtime::{
-    Config, Engine, Func, Global, GlobalType, Instance, Linker, Module, Mutability, Result, Store,
-    Val, ValType,
+    Config, Engine, Func, Global, GlobalType, Instance, InstanceAllocationStrategy, Linker, Module,
+    Mutability, PoolingAllocationConfig, Result, Store, Val, ValType,
 };
 use wasmtime_fuzzing::oom::OomTest;
 
@@ -326,6 +326,32 @@ fn instance_global_init_with_imported_global() -> Result<()> {
     })
 }
 
+#[test]
+fn pooling_allocator_memory_slot_leak_on_oom() -> Result<()> {
+    let mut pool = PoolingAllocationConfig::default();
+    pool.total_memories(1);
+    pool.total_tables(10);
+    pool.total_core_instances(10);
+
+    let mut config = Config::new();
+    config.concurrency_support(false);
+    config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool));
+    let engine = Engine::new(&config)?;
+
+    let module = Module::new(&engine, r#"(module (memory (export "m") 1))"#)?;
+    let linker = Linker::<()>::new(&engine);
+    let instance_pre = linker.instantiate_pre(&module)?;
+
+    // With only one memory slot in the pooling allocator shared amongst all
+    // runs here we should be guaranteed that an OOM failure anywhere always
+    // returns the memory, if any, back to the pooling allocator.
+    OomTest::new().allow_alloc_after_oom(true).test(|| {
+        let mut store = Store::try_new(&engine, ())?;
+        let _instance = instance_pre.instantiate(&mut store)?;
+        Ok(())
+    })
+}
+
 #[test]
 fn table_element_segment_init() -> Result<()> {
     let module_bytes = {

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

@@ -151,10 +151,14 @@ pub struct Instance {
 }
 
 impl Instance {
-    /// Create an instance at the given memory address.
+    /// Creates a new owned instance handle from `req`.
     ///
-    /// It is assumed the memory was properly aligned and the
-    /// allocation was `alloc_size` in bytes.
+    /// The runtime memory/table data structures must have been previously
+    /// allocated and are present within `memories` and `tables`. These values
+    /// are `mem::take`n upon allocation success of an instance, and if the
+    /// instance allocation fails then these are otherwise left in place. This
+    /// enable the pooling allocator to run custom deallocation code for them,
+    /// for example.
     ///
     /// # Safety
     ///
@@ -163,8 +167,8 @@ impl Instance {
     /// and `tables` must have been allocated for `req.store`.
     unsafe fn new(
         req: InstanceAllocationRequest,
-        memories: TryPrimaryMap<DefinedMemoryIndex, (MemoryAllocationIndex, Memory)>,
-        tables: TryPrimaryMap<DefinedTableIndex, (TableAllocationIndex, Table)>,
+        memories: &mut TryPrimaryMap<DefinedMemoryIndex, (MemoryAllocationIndex, Memory)>,
+        tables: &mut TryPrimaryMap<DefinedTableIndex, (TableAllocationIndex, Table)>,
     ) -> Result<InstanceHandle, OutOfMemory> {
         let module = req.runtime_info.env_module();
         let memory_tys = &module.memories;
@@ -191,18 +195,28 @@ impl Instance {
             passive_elements.push(PassiveElementSegment::new(*ty, len)?)?;
         }
 
+        // Allocate the instance and its `VMContext` with empty memory and table
+        // maps. This is the final fallible allocation in this function; only
+        // after it succeeds do we transfer ownership of the pool-allocated
+        // `memories`/`tables` into the instance (below), so that a failure here
+        // leaves them in the caller's deallocation guard to be freed.
         let mut ret = OwnedInstance::new(Instance {
             id: req.id,
             runtime_info: req.runtime_info.clone(),
-            memories,
-            tables,
+            memories: TryPrimaryMap::default(),
+            tables: TryPrimaryMap::default(),
             passive_elements,
             #[cfg(feature = "wmemcheck")]
             wmemcheck_state,
             store: None,
             vmctx: OwnedVMContext::new(),
         })?;
 
+        // Can't fail any more, so transfer ownership of `memories` and `tables`
+        // to this instance.
+        *ret.get_mut().memories_mut() = mem::take(memories);
+        *ret.get_mut().tables_mut() = mem::take(tables);
+
         // SAFETY: this vmctx was allocated with the same layout above, so it
         // should be safe to initialize with the same values here.
         unsafe {

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

@@ -324,32 +324,35 @@ impl dyn InstanceAllocator + '_ {
                 .expect("module should have already been validated before allocation");
         }
 
-        self.increment_core_instance_count()?;
-
         let num_defined_memories = module.num_defined_memories();
         let num_defined_tables = module.num_defined_tables();
 
+        let memories = TryPrimaryMap::with_capacity(num_defined_memories)?;
+        let tables = TryPrimaryMap::with_capacity(num_defined_tables)?;
+
+        // Note that incrementing the instance count here must be done just
+        // before creation of `DeallocateOnDrop`. This is required to ensure
+        // that upon successful increment it'll get paired with a decrement
+        // below should anything fail.
+        self.increment_core_instance_count()?;
         let mut guard = DeallocateOnDrop {
             run_deallocate: true,
-            memories: TryPrimaryMap::with_capacity(num_defined_memories)?,
-            tables: TryPrimaryMap::with_capacity(num_defined_tables)?,
+            memories,
+            tables,
             allocator: self,
+            // NB: do not add more initialization here if it can fail, move that
+            // above the increment above instead.
         };
 
         self.allocate_memories(&mut request, &mut guard.memories)
             .await?;
         self.allocate_tables(&mut request, &mut guard.tables)
             .await?;
+
         // SAFETY: memories/tables were just allocated from the store within
         // `request` and this function's own contract requires that the
         // imports are valid.
-        let handle = unsafe {
-            Instance::new(
-                request,
-                mem::take(&mut guard.memories),
-                mem::take(&mut guard.tables),
-            )?
-        };
+        let handle = unsafe { Instance::new(request, &mut guard.memories, &mut guard.tables)? };
         guard.run_deallocate = false;
         return Ok(handle);
 
@@ -363,6 +366,8 @@ impl dyn InstanceAllocator + '_ {
         impl Drop for DeallocateOnDrop<'_> {
             fn drop(&mut self) {
                 if !self.run_deallocate {
+                    debug_assert!(self.memories.is_empty());
+                    debug_assert!(self.tables.is_empty());
                     return;
                 }
                 // SAFETY: these were previously allocated by this allocator