Refactor resetting memory on MemoryImageSlot drop (#11510) (#11537)

* Refactor resetting memory on `MemoryImageSlot` drop

This commit refactors the behavior of dropping a `MemoryImageSlot` to no
longer map anonymous memory into the slot. This behavior was implemented
previously because if a `MemoryImageSlot` is dropped then the state of
the slot is unknown and to prevent any sort of data leakage a reset is
performed.

This reset operation, however, is fallible in that it calls `mmap`.
Calls to `mmap` can fail due to `ENOMEM`, for example, if the process
has reached its VMA limit. This means that if a process is in a near-OOM
condition then failing to allocate a memory image could panic the
process due to the `unwrap()` in the destructor of `MemoryImageSlot`.
The purpose of this commit is to avoid this `unwrap()` and instead move
the reset behavior to a location where an error can be propagated.

This commit removes the clear-on-drop behavior of `MemoryImageSlot`
slot. This was already disabled everywhere except the pooling allocator.
The pooling allocator now maintains an extra bit of state-per-slot where
instead of storing `Option<MemoryImageSlot>` it now stores effectively
one other variant of "unknown". On reuse of an "unknown" slot the memory
is reset back to an anonymous mapping and this is all done in a context
where an error can be propagated.

Two tests are added in this commit to confirm all of this behavior:

* The first test is a new test that passes both before and after this
  commit which performs a failed allocation of a memory slot. A
  successful allocation is then made to ensure that the previous image
  is not present and zero memory is present. This test fails before the
  commit if the clear-on-drop behavior is removed, and it fails with this
  commit if the clear-on-reusing-unknown behavior is removed.
  Effectively this test ensures that the clear-on-unknown-state logic is
  present.

* The second test is a new test that panicked before this commit and
  passes afterwards. This second test exhausts all VMAs in the current
  process, or at least most of them, and then tries to allocate some
  instances with an image. Instance allocation will eventually fail and
  cause the erroneous path to get executed. This previously unwrapped a
  `ENOMEM` failure, and now it can be handled gracefully by the
  embedder.

* Skip the new test on QEMU, it fails on CI

* Only run test on Linux

Author: alexcrichton

RELEASES.md

@@ -1,3 +1,15 @@
+## 36.0.2
+
+Released 2025-08-26.
+
+### Fixed
+
+* Wasmtime will no longer panic in the pooling allocator when in near-OOM
+  conditions related to resetting the linear memory of a slot.
+  [#11510](https://github.com/bytecodealliance/wasmtime/pull/11510)
+
+--------------------------------------------------------------------------------
+
 ## 36.0.1
 
 Released 2025-08-21.

crates/wasmtime/Cargo.toml

@@ -113,6 +113,10 @@ wasmtime-versioned-export-macros = { workspace = true, optional = true }
 name = "host_segfault"
 harness = false
 
+[[test]]
+name = "pooling_alloc_near_oom"
+harness = false
+
 # =============================================================================
 #
 # Features for the Wasmtime crate.

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

@@ -317,12 +317,6 @@ pub struct MemoryImageSlot {
     /// initial image content, as appropriate. Everything between
     /// `self.accessible` and `self.static_size` is inaccessible.
     dirty: bool,
-
-    /// Whether this MemoryImageSlot is responsible for mapping anonymous
-    /// memory (to hold the reservation while overwriting mappings
-    /// specific to this slot) in place when it is dropped. Default
-    /// on, unless the caller knows what they are doing.
-    clear_on_drop: bool,
 }
 
 impl MemoryImageSlot {
@@ -345,18 +339,9 @@ impl MemoryImageSlot {
             accessible,
             image: None,
             dirty: false,
-            clear_on_drop: true,
         }
     }
 
-    /// Inform the MemoryImageSlot that it should *not* clear the underlying
-    /// address space when dropped. This should be used only when the
-    /// caller will clear or reuse the address space in some other
-    /// way.
-    pub(crate) fn no_clear_on_drop(&mut self) {
-        self.clear_on_drop = false;
-    }
-
     pub(crate) fn set_heap_limit(&mut self, size_bytes: usize) -> Result<()> {
         let size_bytes_aligned = HostAlignedByteCount::new_rounded_up(size_bytes)?;
         assert!(size_bytes <= self.static_size);
@@ -747,7 +732,7 @@ impl MemoryImageSlot {
 
     /// Map anonymous zeroed memory across the whole slot,
     /// inaccessible. Used both during instantiate and during drop.
-    fn reset_with_anon_memory(&mut self) -> Result<()> {
+    pub(crate) fn reset_with_anon_memory(&mut self) -> Result<()> {
         if self.static_size == 0 {
             assert!(self.image.is_none());
             assert_eq!(self.accessible, 0);
@@ -765,45 +750,6 @@ impl MemoryImageSlot {
     }
 }
 
-impl Drop for MemoryImageSlot {
-    fn drop(&mut self) {
-        // The MemoryImageSlot may be dropped if there is an error during
-        // instantiation: for example, if a memory-growth limiter
-        // disallows a guest from having a memory of a certain size,
-        // after we've already initialized the MemoryImageSlot.
-        //
-        // We need to return this region of the large pool mmap to a
-        // safe state (with no module-specific mappings). The
-        // MemoryImageSlot will not be returned to the MemoryPool, so a new
-        // MemoryImageSlot will be created and overwrite the mappings anyway
-        // on the slot's next use; but for safety and to avoid
-        // resource leaks it's better not to have stale mappings to a
-        // possibly-otherwise-dead module's image.
-        //
-        // To "wipe the slate clean", let's do a mmap of anonymous
-        // memory over the whole region, with PROT_NONE. Note that we
-        // *can't* simply munmap, because that leaves a hole in the
-        // middle of the pooling allocator's big memory area that some
-        // other random mmap may swoop in and take, to be trampled
-        // over by the next MemoryImageSlot later.
-        //
-        // Since we're in drop(), we can't sanely return an error if
-        // this mmap fails. Instead though the result is unwrapped here to
-        // trigger a panic if something goes wrong. Otherwise if this
-        // reset-the-mapping fails then on reuse it might be possible, depending
-        // on precisely where errors happened, that stale memory could get
-        // leaked through.
-        //
-        // The exception to all of this is if the `clear_on_drop` flag
-        // (which is set by default) is false. If so, the owner of
-        // this MemoryImageSlot has indicated that it will clean up in some
-        // other way.
-        if self.clear_on_drop {
-            self.reset_with_anon_memory().unwrap();
-        }
-    }
-}
-
 #[cfg(all(test, target_os = "linux", not(miri)))]
 mod test {
     use super::*;
@@ -878,7 +824,6 @@ mod test {
         // Create a MemoryImageSlot on top of it
         let mut memfd =
             MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
-        memfd.no_clear_on_drop();
         assert!(!memfd.is_dirty());
         // instantiate with 64 KiB initial size
         memfd.instantiate(64 << 10, None, &ty, &tunables).unwrap();
@@ -926,7 +871,6 @@ mod test {
         // Create a MemoryImageSlot on top of it
         let mut memfd =
             MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
-        memfd.no_clear_on_drop();
         // Create an image with some data.
         let image = Arc::new(create_memfd_with_data(page_size, &[1, 2, 3, 4]).unwrap());
         // Instantiate with this image
@@ -1016,7 +960,6 @@ mod test {
         let mmap = mmap_4mib_inaccessible();
         let mut memfd =
             MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
-        memfd.no_clear_on_drop();
 
         // Test basics with the image
         for image_off in [0, page_size, page_size * 2] {
@@ -1083,7 +1026,6 @@ mod test {
         let mmap = mmap_4mib_inaccessible();
         let mut memfd =
             MemoryImageSlot::create(mmap.zero_offset(), HostAlignedByteCount::ZERO, 4 << 20);
-        memfd.no_clear_on_drop();
         let image = Arc::new(create_memfd_with_data(page_size, &[1, 2, 3, 4]).unwrap());
         let initial = 64 << 10;
 

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

@@ -64,6 +64,7 @@ use crate::{
     runtime::vm::mpk::{self, ProtectionKey, ProtectionMask},
     vm::HostAlignedByteCount,
 };
+use std::mem;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
 use wasmtime_environ::{DefinedMemoryIndex, Module, Tunables};
@@ -112,7 +113,7 @@ pub struct MemoryPool {
 
     /// If using a copy-on-write allocation scheme, the slot management. We
     /// dynamically transfer ownership of a slot to a Memory when in use.
-    image_slots: Vec<Mutex<Option<MemoryImageSlot>>>,
+    image_slots: Vec<Mutex<ImageSlot>>,
 
     /// A description of the various memory sizes used in allocating the
     /// `mapping` slab.
@@ -139,6 +140,37 @@ pub struct MemoryPool {
     next_available_pkey: AtomicUsize,
 }
 
+/// The state of memory for each slot in this pool.
+#[derive(Debug)]
+enum ImageSlot {
+    /// This slot is guaranteed to be entirely unmapped.
+    ///
+    /// This is the initial state of all slots.
+    Unmapped,
+
+    /// The state of this slot is unknown.
+    ///
+    /// This encompasses a number of situations such as:
+    ///
+    /// * The slot is currently in use.
+    /// * The slot was attempted to be in use, but allocation failed.
+    /// * The slot was used but not deallocated properly.
+    ///
+    /// All of these situations are lumped into this one variant indicating
+    /// that, at a base level, no knowledge is known about this slot. Using a
+    /// slot in this state first requires resetting all memory in this slot by
+    /// mapping anonymous memory on top of the entire slot.
+    Unknown,
+
+    /// This slot was previously used and `MemoryImageSlot` maintains the state
+    /// about what this slot was last configured as.
+    ///
+    /// Future use of this slot will use `MemoryImageSlot` to continue to
+    /// re-instantiate and reuse images and such. This state is entered after
+    /// and allocated slot is successfully deallcoated.
+    PreviouslyUsed(MemoryImageSlot),
+}
+
 impl MemoryPool {
     /// Create a new `MemoryPool`.
     pub fn new(config: &PoolingInstanceAllocatorConfig, tunables: &Tunables) -> Result<Self> {
@@ -219,7 +251,7 @@ impl MemoryPool {
             );
         }
 
-        let image_slots: Vec<_> = std::iter::repeat_with(|| Mutex::new(None))
+        let image_slots: Vec<_> = std::iter::repeat_with(|| Mutex::new(ImageSlot::Unmapped))
             .take(constraints.num_slots)
             .collect();
 
@@ -360,7 +392,7 @@ impl MemoryPool {
             let base = self.get_base(allocation_index);
             let base_capacity = self.layout.max_memory_bytes;
 
-            let mut slot = self.take_memory_image_slot(allocation_index);
+            let mut slot = self.take_memory_image_slot(allocation_index)?;
             let image = match memory_index {
                 Some(memory_index) => request.runtime_info.memory_image(memory_index)?,
                 None => None,
@@ -459,13 +491,18 @@ impl MemoryPool {
                     .allocator
                     .alloc_affine_and_clear_affinity(module, memory_index)
                 {
-                    // Clear the image from the slot and, if successful, return it back
-                    // to our state. Note that on failure here the whole slot will get
-                    // paved over with an anonymous mapping.
+                    // Attempt to acquire the `MemoryImageSlot` state for this
+                    // slot, and then if we have that try to remove the image,
+                    // and then if all that succeeds put the slot back in.
+                    //
+                    // If anything fails then the slot will be in an "unknown"
+                    // state which means that on next use it'll be remapped with
+                    // anonymous memory.
                     let index = MemoryAllocationIndex(id.0);
-                    let mut slot = self.take_memory_image_slot(index);
-                    if slot.remove_image().is_ok() {
-                        self.return_memory_image_slot(index, slot);
+                    if let Ok(mut slot) = self.take_memory_image_slot(index) {
+                        if slot.remove_image().is_ok() {
+                            self.return_memory_image_slot(index, slot);
+                        }
                     }
 
                     stripe.allocator.free(id);
@@ -485,21 +522,49 @@ impl MemoryPool {
         self.mapping.offset(offset).expect("offset is in bounds")
     }
 
-    /// Take ownership of the given image slot. Must be returned via
-    /// `return_memory_image_slot` when the instance is done using it.
-    fn take_memory_image_slot(&self, allocation_index: MemoryAllocationIndex) -> MemoryImageSlot {
-        let maybe_slot = self.image_slots[allocation_index.index()]
-            .lock()
-            .unwrap()
-            .take();
-
-        maybe_slot.unwrap_or_else(|| {
+    /// Take ownership of the given image slot.
+    ///
+    /// This method is used when a `MemoryAllocationIndex` has been allocated
+    /// and the state of the slot needs to be acquired. This will lazily
+    /// allocate a `MemoryImageSlot` which describes the current (and possibly
+    /// prior) state of the slot.
+    ///
+    /// During deallocation this structure is passed back to
+    /// `return_memory_image_slot`.
+    ///
+    /// Note that this is a fallible method because using a slot might require
+    /// resetting the memory that was previously there. This reset operation
+    /// is a fallible operation that may not succeed. If it fails then this
+    /// slot cannot be used at this time.
+    fn take_memory_image_slot(
+        &self,
+        allocation_index: MemoryAllocationIndex,
+    ) -> Result<MemoryImageSlot> {
+        let (maybe_slot, needs_reset) = {
+            let mut slot = self.image_slots[allocation_index.index()].lock().unwrap();
+            match mem::replace(&mut *slot, ImageSlot::Unknown) {
+                ImageSlot::Unmapped => (None, false),
+                ImageSlot::Unknown => (None, true),
+                ImageSlot::PreviouslyUsed(state) => (Some(state), false),
+            }
+        };
+        let mut slot = maybe_slot.unwrap_or_else(|| {
             MemoryImageSlot::create(
                 self.get_base(allocation_index),
                 HostAlignedByteCount::ZERO,
                 self.layout.max_memory_bytes.byte_count(),
             )
-        })
+        });
+
+        // For `Unknown` slots it means that `slot` is brand new and isn't
+        // actually tracking the state of the previous slot, so reset it
+        // entirely with anonymous memory to wipe the slate clean and start
+        // from zero. This should only happen if allocation of the previous
+        // slot failed, for example.
+        if needs_reset {
+            slot.reset_with_anon_memory()?;
+        }
+        Ok(slot)
     }
 
     /// Return ownership of the given image slot.
@@ -509,21 +574,12 @@ impl MemoryPool {
         slot: MemoryImageSlot,
     ) {
         assert!(!slot.is_dirty());
-        *self.image_slots[allocation_index.index()].lock().unwrap() = Some(slot);
-    }
-}
 
-impl Drop for MemoryPool {
-    fn drop(&mut self) {
-        // Clear the `clear_no_drop` flag (i.e., ask to *not* clear on
-        // drop) for all slots, and then drop them here. This is
-        // valid because the one `Mmap` that covers the whole region
-        // can just do its one munmap.
-        for mut slot in std::mem::take(&mut self.image_slots) {
-            if let Some(slot) = slot.get_mut().unwrap() {
-                slot.no_clear_on_drop();
-            }
-        }
+        let prev = mem::replace(
+            &mut *self.image_slots[allocation_index.index()].lock().unwrap(),
+            ImageSlot::PreviouslyUsed(slot),
+        );
+        assert!(matches!(prev, ImageSlot::Unknown));
     }
 }
 

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

@@ -547,29 +547,6 @@ impl LocalMemory {
 
                     let mut slot =
                         MemoryImageSlot::create(mmap_base, byte_size, alloc.byte_capacity());
-                    // On drop, we will unmap our mmap'd range that this slot
-                    // was mapped on top of, so there is no need for the slot to
-                    // wipe it with an anonymous mapping first.
-                    //
-                    // Note that this code would be incorrect if clear-on-drop
-                    // were enabled. That's because:
-                    //
-                    // * In the struct definition, `memory_image` above is listed
-                    //   after `alloc`.
-                    // * Rust drops fields in the order they're defined, so
-                    //   `memory_image` would be dropped after `alloc`.
-                    // * `alloc` can represent either owned memory (i.e. the mmap is
-                    //   freed on drop) or logically borrowed memory (something else
-                    //   manages the mmap).
-                    // * If `alloc` is borrowed memory, then this isn't an issue.
-                    // * But if `alloc` is owned memory, then it would first drop
-                    //   the mmap, and then `memory_image` would try to remap
-                    //   part of that same memory as part of clear-on-drop.
-                    //
-                    // A lot of this really suggests representing the ownership
-                    // via Rust lifetimes -- that would be a major refactor,
-                    // though.
-                    slot.no_clear_on_drop();
                     slot.instantiate(alloc.byte_size(), Some(image), ty, tunables)?;
                     Some(slot)
                 } else {