Clear affine slots when dropping a Module

crates/runtime/src/cow.rs

@@ -499,14 +499,8 @@ impl MemoryImageSlot {
         // extent of the prior initialization image in order to preserve
         // resident memory that might come before or after the image.
         if self.image.as_ref() != maybe_image {
-            if let Some(image) = &self.image {
-                unsafe {
-                    image
-                        .remap_as_zeros_at(self.base)
-                        .map_err(|e| InstantiationError::Resource(e.into()))?;
-                }
-                self.image = None;
-            }
+            self.remove_image()
+                .map_err(|e| InstantiationError::Resource(e.into()))?;
         }
 
         // The next order of business is to ensure that `self.accessible` is
@@ -565,6 +559,16 @@ impl MemoryImageSlot {
         Ok(())
     }
 
+    pub(crate) fn remove_image(&mut self) -> Result<()> {
+        if let Some(image) = &self.image {
+            unsafe {
+                image.remap_as_zeros_at(self.base)?;
+            }
+            self.image = None;
+        }
+        Ok(())
+    }
+
     /// Resets this linear memory slot back to a "pristine state".
     ///
     /// This will reset the memory back to its original contents on Linux or

crates/runtime/src/instance/allocator.rs

@@ -2,8 +2,7 @@ use crate::imports::Imports;
 use crate::instance::{Instance, InstanceHandle, RuntimeMemoryCreator};
 use crate::memory::{DefaultMemoryCreator, Memory};
 use crate::table::Table;
-use crate::ModuleRuntimeInfo;
-use crate::Store;
+use crate::{CompiledModuleId, ModuleRuntimeInfo, Store};
 use anyhow::Result;
 use std::alloc;
 use std::any::Any;
@@ -190,6 +189,13 @@ pub unsafe trait InstanceAllocator: Send + Sync {
     /// The provided stack is required to have been allocated with `allocate_fiber_stack`.
     #[cfg(feature = "async")]
     unsafe fn deallocate_fiber_stack(&self, stack: &wasmtime_fiber::FiberStack);
+
+    /// Purges all lingering resources related to `module` from within this
+    /// allocator.
+    ///
+    /// Primarily present for the pooling allocator to remove mappings of
+    /// this module from slots in linear memory.
+    fn purge_module(&self, module: CompiledModuleId);
 }
 
 fn get_table_init_start(
@@ -593,4 +599,6 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
     unsafe fn deallocate_fiber_stack(&self, _stack: &wasmtime_fiber::FiberStack) {
         // The on-demand allocator has no further bookkeeping for fiber stacks
     }
+
+    fn purge_module(&self, _: CompiledModuleId) {}
 }

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

@@ -12,7 +12,7 @@ use super::{
     InstantiationError,
 };
 use crate::{instance::Instance, Memory, Mmap, Table};
-use crate::{MemoryImageSlot, ModuleRuntimeInfo, Store};
+use crate::{CompiledModuleId, MemoryImageSlot, ModuleRuntimeInfo, Store};
 use anyhow::{anyhow, bail, Context, Result};
 use libc::c_void;
 use std::convert::TryFrom;
@@ -560,6 +560,22 @@ impl InstancePool {
 
         bail!("{}", message)
     }
+
+    fn purge_module(&self, module: CompiledModuleId) {
+        // Purging everything related to `module` primarily means clearing out
+        // all of its memory images present in the virtual address space. Go
+        // through the index allocator for slots affine to `module` and reset
+        // them, freeing up the index when we're done.
+        //
+        // Note that this is only called when the specified `module` won't be
+        // allocated further (the module is being dropped) so this shouldn't hit
+        // any sort of infinite loop since this should be the final operation
+        // working with `module`.
+        while let Some(index) = self.index_allocator.alloc_affine_and_clear_affinity(module) {
+            self.memories.clear_images(index.0);
+            self.index_allocator.free(index);
+        }
+    }
 }
 
 /// Represents a pool of WebAssembly linear memories.
@@ -740,6 +756,26 @@ impl MemoryPool {
         let idx = instance_index * self.max_memories + (memory_index.as_u32() as usize);
         *self.image_slots[idx].lock().unwrap() = Some(slot);
     }
+
+    /// Resets all the images for the instance index slot specified to clear out
+    /// any prior mappings.
+    ///
+    /// This is used when a `Module` is dropped at the `wasmtime` layer to clear
+    /// out any remaining mappings and ensure that its memfd backing, if any, is
+    /// removed from the address space to avoid lingering references to it.
+    fn clear_images(&self, instance_index: usize) {
+        for i in 0..self.max_memories {
+            let index = DefinedMemoryIndex::from_u32(i as u32);
+
+            // 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.
+            let mut slot = self.take_memory_image_slot(instance_index, index);
+            if slot.remove_image().is_ok() {
+                self.return_memory_image_slot(instance_index, index, slot);
+            }
+        }
+    }
 }
 
 impl Drop for MemoryPool {
@@ -1116,6 +1152,10 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator {
     unsafe fn deallocate_fiber_stack(&self, _stack: &wasmtime_fiber::FiberStack) {
         // A no-op as we don't own the fiber stack on Windows
     }
+
+    fn purge_module(&self, module: CompiledModuleId) {
+        self.instances.purge_module(module);
+    }
 }
 
 #[cfg(test)]