Fix leaking GC heap slots on errors

crates/wasmtime/src/runtime/store.rs

@@ -1846,6 +1846,9 @@ impl StoreOpaque {
                 engine.features().gc_types(),
                 "cannot allocate a GC store when GC is disabled at configuration time"
             );
+            let gc_runtime = engine
+                .gc_runtime()
+                .context("no GC runtime: GC disabled at compile time or configuration time")?;
 
             // First, allocate the memory that will be our GC heap's storage.
             let mut request = InstanceAllocationRequest {
@@ -1868,13 +1871,20 @@ impl StoreOpaque {
 
             // Then, allocate the actual GC heap, passing in that memory
             // storage.
-            let gc_runtime = engine
-                .gc_runtime()
-                .context("no GC runtime: GC disabled at compile time or configuration time")?;
-            let (index, heap) =
-                engine
+            let (index, mut heap) =
+                match engine
                     .allocator()
-                    .allocate_gc_heap(engine, &**gc_runtime, mem_alloc_index, mem)?;
+                    .allocate_gc_heap(engine, &**gc_runtime, mem_alloc_index)
+                {
+                    Ok(pair) => pair,
+                    Err(e) => unsafe {
+                        engine
+                            .allocator()
+                            .deallocate_memory(None, mem_alloc_index, mem);
+                        return Err(e);
+                    },
+                };
+            heap.attach(mem);
 
             let mut gc_store = GcStore::new(index, heap, engine.tunables().gc_zeal_alloc_counter);
 
@@ -2573,11 +2583,12 @@ impl Drop for StoreOpaque {
             let store_id = self.id();
 
             #[cfg(feature = "gc")]
-            if let Some(gc_store) = self.gc_store.take() {
+            if let Some(mut gc_store) = self.gc_store.take() {
                 let gc_alloc_index = gc_store.allocation_index;
                 log::trace!("store {store_id:?} is deallocating GC heap {gc_alloc_index:?}");
                 debug_assert!(self.engine.features().gc_types());
-                let (mem_alloc_index, mem) =
+                let mem = gc_store.gc_heap.detach();
+                let mem_alloc_index =
                     allocator.deallocate_gc_heap(gc_alloc_index, gc_store.gc_heap);
                 allocator.deallocate_memory(None, mem_alloc_index, mem);
             }

crates/wasmtime/src/runtime/trampoline/memory.rs

@@ -263,7 +263,6 @@ unsafe impl InstanceAllocator for SingleMemoryInstance<'_> {
         _engine: &crate::Engine,
         _gc_runtime: &dyn crate::vm::GcRuntime,
         _memory_alloc_index: crate::vm::MemoryAllocationIndex,
-        _memory: Memory,
     ) -> Result<(crate::vm::GcHeapAllocationIndex, Box<dyn crate::vm::GcHeap>)> {
         unreachable!()
     }
@@ -273,7 +272,7 @@ unsafe impl InstanceAllocator for SingleMemoryInstance<'_> {
         &self,
         _allocation_index: crate::vm::GcHeapAllocationIndex,
         _gc_heap: Box<dyn crate::vm::GcHeap>,
-    ) -> (crate::vm::MemoryAllocationIndex, crate::vm::Memory) {
+    ) -> crate::vm::MemoryAllocationIndex {
         unreachable!()
     }
 }

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

@@ -253,7 +253,6 @@ pub unsafe trait InstanceAllocator: Send + Sync {
         engine: &crate::Engine,
         gc_runtime: &dyn GcRuntime,
         memory_alloc_index: MemoryAllocationIndex,
-        memory: Memory,
     ) -> Result<(GcHeapAllocationIndex, Box<dyn GcHeap>)>;
 
     /// Deallocate a GC heap that was previously allocated with
@@ -265,7 +264,7 @@ pub unsafe trait InstanceAllocator: Send + Sync {
         &self,
         allocation_index: GcHeapAllocationIndex,
         gc_heap: Box<dyn GcHeap>,
-    ) -> (MemoryAllocationIndex, Memory);
+    ) -> MemoryAllocationIndex;
 
     /// Purges all lingering resources related to `module` from within this
     /// allocator.

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

@@ -236,21 +236,20 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
         engine: &crate::Engine,
         gc_runtime: &dyn GcRuntime,
         memory_alloc_index: MemoryAllocationIndex,
-        memory: Memory,
     ) -> Result<(GcHeapAllocationIndex, Box<dyn GcHeap>)> {
         debug_assert_eq!(memory_alloc_index, MemoryAllocationIndex::default());
-        let mut heap = gc_runtime.new_gc_heap(engine)?;
-        heap.attach(memory);
+        let heap = gc_runtime.new_gc_heap(engine)?;
         Ok((GcHeapAllocationIndex::default(), heap))
     }
 
     #[cfg(feature = "gc")]
     fn deallocate_gc_heap(
         &self,
         allocation_index: GcHeapAllocationIndex,
-        mut gc_heap: Box<dyn crate::runtime::vm::GcHeap>,
-    ) -> (MemoryAllocationIndex, Memory) {
+        gc_heap: Box<dyn crate::runtime::vm::GcHeap>,
+    ) -> MemoryAllocationIndex {
         debug_assert_eq!(allocation_index, GcHeapAllocationIndex::default());
-        (MemoryAllocationIndex::default(), gc_heap.detach())
+        debug_assert!(!gc_heap.is_attached());
+        MemoryAllocationIndex::default()
     }
 }

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

@@ -885,14 +885,12 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator {
         engine: &crate::Engine,
         gc_runtime: &dyn GcRuntime,
         memory_alloc_index: MemoryAllocationIndex,
-        memory: Memory,
     ) -> Result<(GcHeapAllocationIndex, Box<dyn GcHeap>)> {
-        let ret = self.gc_heaps.as_ref().unwrap().allocate(
-            engine,
-            gc_runtime,
-            memory_alloc_index,
-            memory,
-        )?;
+        let ret =
+            self.gc_heaps
+                .as_ref()
+                .unwrap()
+                .allocate(engine, gc_runtime, memory_alloc_index)?;
         self.live_gc_heaps.fetch_add(1, Ordering::Relaxed);
         Ok(ret)
     }
@@ -902,7 +900,7 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator {
         &self,
         allocation_index: GcHeapAllocationIndex,
         gc_heap: Box<dyn GcHeap>,
-    ) -> (MemoryAllocationIndex, Memory) {
+    ) -> MemoryAllocationIndex {
         let gc_heaps = self.gc_heaps.as_ref().unwrap();
         self.live_gc_heaps.fetch_sub(1, Ordering::Relaxed);
         gc_heaps.deallocate(allocation_index, gc_heap)

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

@@ -1,7 +1,7 @@
-use super::GcHeapAllocationIndex;
 use super::index_allocator::{SimpleIndexAllocator, SlotId};
+use super::{GcHeapAllocationIndex, PoolConcurrencyLimitError};
 use crate::runtime::vm::{GcHeap, GcRuntime, PoolingInstanceAllocatorConfig, Result};
-use crate::vm::{Memory, MemoryAllocationIndex};
+use crate::vm::MemoryAllocationIndex;
 use crate::{Engine, prelude::*};
 use std::sync::Mutex;
 use wasmtime_environ::Tunables;
@@ -28,11 +28,11 @@ impl HeapSlot {
         }
     }
 
-    fn dealloc(&mut self, heap: Box<dyn GcHeap>) -> MemoryAllocationIndex {
+    fn dealloc(&mut self, heap: Option<Box<dyn GcHeap>>) -> MemoryAllocationIndex {
         match *self {
             HeapSlot::Free(_) => panic!("already free"),
             HeapSlot::InUse(memory_alloc_index) => {
-                *self = HeapSlot::Free(Some(heap));
+                *self = HeapSlot::Free(heap);
                 memory_alloc_index
             }
         }
@@ -103,33 +103,32 @@ impl GcHeapPool {
         engine: &Engine,
         gc_runtime: &dyn GcRuntime,
         memory_alloc_index: MemoryAllocationIndex,
-        memory: Memory,
     ) -> Result<(GcHeapAllocationIndex, Box<dyn GcHeap>)> {
         let allocation_index = self
             .index_allocator
             .alloc()
             .map(|slot| GcHeapAllocationIndex(slot.0))
-            .ok_or_else(|| {
-                format_err!(
-                    "maximum concurrent GC heap limit of {} reached",
-                    self.max_gc_heaps
-                )
-            })?;
+            .ok_or_else(|| PoolConcurrencyLimitError::new(self.max_gc_heaps, "GC heaps"))?;
         debug_assert_ne!(allocation_index, GcHeapAllocationIndex::default());
 
-        let mut heap = match {
+        let heap = match {
             let mut heaps = self.heaps.lock().unwrap();
             heaps[allocation_index.index()].alloc(memory_alloc_index)
         } {
             // If we already have a heap at this slot, reuse it.
             Some(heap) => heap,
             // Otherwise, we haven't forced this slot's lazily allocated heap
             // yet. So do that now.
-            None => gc_runtime.new_gc_heap(engine)?,
+            None => match gc_runtime.new_gc_heap(engine) {
+                Ok(heap) => heap,
+                Err(e) => {
+                    self.deallocate_maybe_with_heap(allocation_index, None);
+                    return Err(e);
+                }
+            },
         };
 
         debug_assert!(!heap.is_attached());
-        heap.attach(memory);
 
         Ok((allocation_index, heap))
     }
@@ -138,11 +137,20 @@ impl GcHeapPool {
     pub fn deallocate(
         &self,
         allocation_index: GcHeapAllocationIndex,
-        mut heap: Box<dyn GcHeap>,
-    ) -> (MemoryAllocationIndex, Memory) {
-        debug_assert_ne!(allocation_index, GcHeapAllocationIndex::default());
+        heap: Box<dyn GcHeap>,
+    ) -> MemoryAllocationIndex {
+        self.deallocate_maybe_with_heap(allocation_index, Some(heap))
+    }
 
-        let memory = heap.detach();
+    fn deallocate_maybe_with_heap(
+        &self,
+        allocation_index: GcHeapAllocationIndex,
+        heap: Option<Box<dyn GcHeap>>,
+    ) -> MemoryAllocationIndex {
+        debug_assert_ne!(allocation_index, GcHeapAllocationIndex::default());
+        if let Some(heap) = &heap {
+            debug_assert!(!heap.is_attached());
+        }
 
         // NB: Replace the heap before freeing the index. If we did it in the
         // opposite order, a concurrent allocation request could reallocate the
@@ -155,6 +163,6 @@ impl GcHeapPool {
 
         self.index_allocator.free(SlotId(allocation_index.0), 0);
 
-        (memory_alloc_index, memory)
+        memory_alloc_index
     }
 }

tests/all/gc.rs

@@ -3636,3 +3636,56 @@ async fn issue_13516_copying_collector_cancellation_during_gc() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+#[cfg_attr(miri, ignore)]
+fn pooling_gc_heap_failure_does_not_leak_memory_slot() -> Result<()> {
+    if crate::skip_pooling_allocator_tests() {
+        return Ok(());
+    }
+
+    const TOTAL_MEMORIES: u32 = 4;
+
+    let mut pool = crate::small_pool_config();
+    pool.total_memories(TOTAL_MEMORIES)
+        .total_gc_heaps(1)
+        .total_core_instances(TOTAL_MEMORIES + 2)
+        .total_tables(TOTAL_MEMORIES + 2)
+        .memory_protection_keys(Enabled::No);
+
+    let mut config = Config::new();
+    config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool));
+    let engine = Engine::new(&config)?;
+
+    let gc_module = Module::new(
+        &engine,
+        r#"
+            (module
+                (type $s (struct))
+                (func (export "alloc")
+                    (drop (struct.new $s)))
+            )
+        "#,
+    )?;
+
+    // Hold the gc slot, then assert other attempts are failures
+    {
+        let mut held = Store::new(&engine, ());
+        Instance::new(&mut held, &gc_module, &[])?;
+
+        for _ in 0..10 {
+            let mut store = Store::new(&engine, ());
+            let err = Instance::new(&mut store, &gc_module, &[]).unwrap_err();
+            assert!(err.is::<PoolConcurrencyLimitError>(), "bad error: {err:?}");
+        }
+    }
+
+    // should be able to use all slots now...
+    let mut store = Store::new(&engine, ());
+    let mem_module = Module::new(&engine, r#"(module (memory 1 1))"#)?;
+    for _ in 0..TOTAL_MEMORIES {
+        Instance::new(&mut store, &mem_module, &[]).unwrap();
+    }
+
+    Ok(())
+}