mpk: allow benchmarking MPK

benches/call.rs

@@ -23,6 +23,7 @@ enum IsAsync {
     Yes,
     YesPooling,
     No,
+    NoPooling,
 }
 
 impl IsAsync {
@@ -31,12 +32,13 @@ impl IsAsync {
             IsAsync::Yes => "async",
             IsAsync::YesPooling => "async-pool",
             IsAsync::No => "sync",
+            IsAsync::NoPooling => "sync-pool",
         }
     }
     fn use_async(&self) -> bool {
         match self {
             IsAsync::Yes | IsAsync::YesPooling => true,
-            IsAsync::No => false,
+            IsAsync::No | IsAsync::NoPooling => false,
         }
     }
 }
@@ -47,14 +49,29 @@ fn engines() -> Vec<(Engine, IsAsync)> {
     #[cfg(feature = "component-model")]
     config.wasm_component_model(true);
 
+    let mut pool = PoolingAllocationConfig::default();
+    if std::env::var("WASMTIME_TEST_FORCE_MPK").is_ok() {
+        pool.memory_protection_keys(MpkEnabled::Enable);
+    }
+
     vec![
         (Engine::new(&config).unwrap(), IsAsync::No),
+        (
+            Engine::new(
+                config
+                    .clone()
+                    .allocation_strategy(InstanceAllocationStrategy::Pooling(pool.clone())),
+            )
+            .unwrap(),
+            IsAsync::NoPooling,
+        ),
         (
             Engine::new(config.async_support(true)).unwrap(),
             IsAsync::Yes,
         ),
         (
-            Engine::new(config.allocation_strategy(InstanceAllocationStrategy::pooling())).unwrap(),
+            Engine::new(config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool)))
+                .unwrap(),
             IsAsync::YesPooling,
         ),
     ]

crates/runtime/src/mpk/enabled.rs

@@ -50,7 +50,11 @@ static KEYS: OnceLock<Vec<ProtectionKey>> = OnceLock::new();
 /// Any accesses to pages marked by another key will result in a `SIGSEGV`
 /// fault.
 pub fn allow(mask: ProtectionMask) {
-    let previous = pkru::read();
+    let previous = if log::log_enabled!(log::Level::Trace) {
+        pkru::read()
+    } else {
+        0
+    };
     pkru::write(mask.0);
     log::trace!("PKRU change: {:#034b} => {:#034b}", previous, pkru::read());
 }

crates/runtime/src/mpk/pkru.rs

@@ -30,6 +30,7 @@ pub const ALLOW_ACCESS: u32 = 0;
 pub const DISABLE_ACCESS: u32 = 0b11111111_11111111_11111111_11111111;
 
 /// Read the value of the `PKRU` register.
+#[inline]
 pub fn read() -> u32 {
     // ECX must be 0 to prevent a general protection exception (#GP).
     let ecx: u32 = 0;
@@ -42,6 +43,7 @@ pub fn read() -> u32 {
 }
 
 /// Write a value to the `PKRU` register.
+#[inline]
 pub fn write(pkru: u32) {
     // Both ECX and EDX must be 0 to prevent a general protection exception
     // (#GP).