mpk: Fix index used when purging a module in the pooling allocator (#12910)

This commit fixes an issue with the pooling allocator when MPK is
enabled, which is off-by-default at compile time. When a module is
dropped all remaining images are purged from the pooling allocator, but
the purging logic mistakenly used the wrong kind of index during purging
which led to corruption of the pooling allocator itself. This fixes the
logic and adds regression tests showcasing the issue as well.

Author: alexcrichton

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

@@ -507,7 +507,7 @@ impl MemoryPool {
         // associated with a module (not just module and memory). The latter
         // would require care to make sure that its maintenance wouldn't be too
         // expensive for normal allocation/free operations.
-        for stripe in &self.stripes {
+        for (stripe_index, stripe) in self.stripes.iter().enumerate() {
             for i in 0..self.memories_per_instance {
                 use wasmtime_environ::EntityRef;
                 let memory_index = DefinedMemoryIndex::new(i);
@@ -522,7 +522,8 @@ impl MemoryPool {
                     // 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 index = StripedAllocationIndex(id.0)
+                        .as_unstriped_slot_index(stripe_index, self.stripes.len());
                     if let Ok(mut slot) = self.take_memory_image_slot(index) {
                         if slot.remove_image().is_ok() {
                             self.return_memory_image_slot(index, Some(slot));

tests/all/pooling_allocator.rs

@@ -1456,3 +1456,41 @@ fn memory_reset_if_instantiation_fails() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn purge_module_with_mpk() -> Result<()> {
+    if !wasmtime::PoolingAllocationConfig::are_memory_protection_keys_available() {
+        println!("skipping test; mpk is not supported");
+        return Ok(());
+    }
+
+    let mut pool = crate::small_pool_config();
+    pool.total_memories(2)
+        .memory_protection_keys(Enabled::Yes)
+        .max_memory_protection_keys(2);
+    let mut config = Config::new();
+    config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool));
+    let engine = Engine::new(&config)?;
+
+    // Create a module and instantiate it in stripe 0.
+    let m1 = Module::new(&engine, "(module (memory 1))")?;
+    let mut store = Store::new(&engine, ());
+    Instance::new(&mut store, &m1, &[])?;
+
+    // Create and instantiate a module in stripe 1. Note that the store is
+    // immediately destroyed here after instantiation. That should leave the
+    // linear memory slot available for re-instantiation.
+    {
+        let m2 = Module::new(&engine, "(module (memory 1))")?;
+        Instance::new(&mut Store::new(&engine, ()), &m2, &[])?;
+
+        // ... drop `m2` here which will purge it and should remove the warm
+        // slot left for the module. Nothing should get corrupted...
+    }
+
+    // Drop the outer store here which will clean up the rest of the
+    // instance/etc.
+    drop(store);
+
+    Ok(())
+}

tests/misc_testsuite/pooling-drop-out-of-order.wast

@@ -0,0 +1,16 @@
+;; Small test case to create a module in one thread, drop another module in
+;; another thread, then drop the first module in the original thread. This
+;; historically exposed an issue with MPK and using striped indices correctly
+;; when purging modules from the pooling allocator.
+
+(module
+  (memory 1)
+  (data (i32.const 0) "\2a\00\00\00")
+  (func (export "load") (result i32)
+    i32.const 0
+    i32.load))
+
+(assert_return (invoke "load") (i32.const 42))
+
+(thread $t1 (module (memory 1)))
+(wait $t1)