From f3b5478bfcb759d99b3910b121c644b4c9c572bf Mon Sep 17 00:00:00 2001
From: Andrew Brown <andrew.brown@intel.com>
Date: Wed, 17 Jan 2024 16:38:45 -0800
Subject: [PATCH] mpk: allow benchmarking MPK (#7787)

* mpk: also force MPK during benchmarking

This change takes advantage of the `WASMTIME_TEST_FORCE_MPK` environment
variable to force its use in the `call.rs` benchmarks. I see differences
in several cases when running:

```console
$ taskset --cpu-list 0-15 cargo bench -- async
$ WASMTIME_TEST_FORCE_MPK=1 taskset --cpu-list 0-15 cargo bench -- async
```

To properly isolate the MPK effects, this adds a `sync-pool` option
(`IsAsync::NoPooling`).

* mpk: disable PKRU read unless logging is turned on

After noticing some minor effects with this PKRU read, I chose to avoid
it unless logging is enabled. A perfectly valid alternative would be to
remove the logging altogether, but I have found this very helpful when
trying to troubleshoot MPK issues.

* mpk: inline PKRU functions

To avoid any unnecessary call overhead, we hint to the compiler that
`pkru::read` and `pkru::write` should be inlined.
---
 benches/call.rs                   | 21 +++++++++++++++++++--
 crates/runtime/src/mpk/enabled.rs |  6 +++++-
 crates/runtime/src/mpk/pkru.rs    |  2 ++
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/benches/call.rs b/benches/call.rs
index 90f9beca60b2..fbfe96d84f3f 100644
--- a/benches/call.rs
+++ b/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,
         ),
     ]
diff --git a/crates/runtime/src/mpk/enabled.rs b/crates/runtime/src/mpk/enabled.rs
index 2ce2de8ea3b4..19827532ced0 100644
--- a/crates/runtime/src/mpk/enabled.rs
+++ b/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());
 }
diff --git a/crates/runtime/src/mpk/pkru.rs b/crates/runtime/src/mpk/pkru.rs
index af6e549466a0..0bb7f63210c7 100644
--- a/crates/runtime/src/mpk/pkru.rs
+++ b/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).