Add a configuration knob for PAGEMAP_SCAN (#11433)

* Add a configuration knob for `PAGEMAP_SCAN`

This commit adds `PoolingAlloationConfig::pagemap_scan` and additionally
adds `-Opooling-pagemap-scan` to configure on the CLI. This is the same
tri-state configuration option as `MpkEnable` so that enum was renamed
to just `Enabled` and repurposed for both options.

This then additionally turns the option off-by-default instead of the
previous on-by-default-if-able-to to enable more slowly rolling out this
feature.

* Fix broken test

Author: alexcrichton

benches/call.rs

@@ -51,7 +51,7 @@ fn engines() -> Vec<(Engine, IsAsync)> {
 
     let mut pool = PoolingAllocationConfig::default();
     if std::env::var("WASMTIME_TEST_FORCE_MPK").is_ok() {
-        pool.memory_protection_keys(MpkEnabled::Enable);
+        pool.memory_protection_keys(Enabled::Yes);
     }
 
     vec![

crates/cli-flags/src/lib.rs

@@ -97,7 +97,7 @@ wasmtime_option_group! {
         /// optimize the size of memory slots.
         #[serde(default)]
         #[serde(deserialize_with = "crate::opt::cli_parse_wrapper")]
-        pub pooling_memory_protection_keys: Option<wasmtime::MpkEnabled>,
+        pub pooling_memory_protection_keys: Option<wasmtime::Enabled>,
 
         /// Sets an upper limit on how many memory protection keys (MPK) Wasmtime
         /// will use. (default: 16)
@@ -193,6 +193,12 @@ wasmtime_option_group! {
 
         /// DEPRECATED: Use `-Cmemory-reservation-for-growth=N` instead.
         pub dynamic_memory_reserved_for_growth: Option<u64>,
+
+        /// Whether or not `PAGEMAP_SCAN` ioctls are used to reset linear
+        /// memory.
+        #[serde(default)]
+        #[serde(deserialize_with = "crate::opt::cli_parse_wrapper")]
+        pub pooling_pagemap_scan: Option<wasmtime::Enabled>,
     }
 
     enum Optimize {
@@ -933,6 +939,9 @@ impl CommonOptions {
                         max => cfg.total_gc_heaps(max),
                         _ => err,
                     }
+                    if let Some(enabled) = self.opts.pooling_pagemap_scan {
+                        cfg.pagemap_scan(enabled);
+                    }
                     config.allocation_strategy(wasmtime::InstanceAllocationStrategy::Pooling(cfg));
                 }
             },

crates/cli-flags/src/opt.rs

@@ -538,22 +538,22 @@ impl WasmtimeOptionValue for wasmtime::Collector {
     }
 }
 
-impl WasmtimeOptionValue for wasmtime::MpkEnabled {
+impl WasmtimeOptionValue for wasmtime::Enabled {
     const VAL_HELP: &'static str = "[=y|n|auto]";
     fn parse(val: Option<&str>) -> Result<Self> {
         match val {
-            None | Some("y") | Some("yes") | Some("true") => Ok(wasmtime::MpkEnabled::Enable),
-            Some("n") | Some("no") | Some("false") => Ok(wasmtime::MpkEnabled::Disable),
-            Some("auto") => Ok(wasmtime::MpkEnabled::Auto),
-            Some(s) => bail!("unknown mpk flag `{s}`, only yes,no,auto,<nothing> accepted"),
+            None | Some("y") | Some("yes") | Some("true") => Ok(wasmtime::Enabled::Yes),
+            Some("n") | Some("no") | Some("false") => Ok(wasmtime::Enabled::No),
+            Some("auto") => Ok(wasmtime::Enabled::Auto),
+            Some(s) => bail!("unknown flag `{s}`, only yes,no,auto,<nothing> accepted"),
         }
     }
 
     fn display(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match *self {
-            wasmtime::MpkEnabled::Enable => f.write_str("y"),
-            wasmtime::MpkEnabled::Disable => f.write_str("n"),
-            wasmtime::MpkEnabled::Auto => f.write_str("auto"),
+            wasmtime::Enabled::Yes => f.write_str("y"),
+            wasmtime::Enabled::No => f.write_str("n"),
+            wasmtime::Enabled::Auto => f.write_str("auto"),
         }
     }
 }

crates/fuzzing/src/generators/config.rs

@@ -5,7 +5,7 @@ use crate::oracles::{StoreLimits, Timeout};
 use anyhow::Result;
 use arbitrary::{Arbitrary, Unstructured};
 use std::time::Duration;
-use wasmtime::{Engine, Module, MpkEnabled, Store};
+use wasmtime::{Enabled, Engine, Module, Store};
 use wasmtime_test_util::wast::{WastConfig, WastTest, limits};
 
 /// Configuration for `wasmtime::Config` and generated modules for a session of
@@ -75,7 +75,7 @@ impl Config {
             pooling.total_memories = config.max_memories as u32;
             pooling.max_memory_size = 10 << 16;
             pooling.max_memories_per_module = config.max_memories as u32;
-            if pooling.memory_protection_keys == MpkEnabled::Auto
+            if pooling.memory_protection_keys == Enabled::Auto
                 && pooling.max_memory_protection_keys > 1
             {
                 pooling.total_memories =

crates/fuzzing/src/generators/pooling_config.rs

@@ -1,7 +1,7 @@
 //! Generate instance limits for the pooling allocation strategy.
 
 use arbitrary::{Arbitrary, Unstructured};
-use wasmtime::MpkEnabled;
+use wasmtime::Enabled;
 
 /// Configuration for `wasmtime::PoolingAllocationStrategy`.
 #[derive(Debug, Clone, Eq, PartialEq, Hash)]
@@ -32,8 +32,10 @@ pub struct PoolingAllocationConfig {
 
     pub async_stack_keep_resident: usize,
 
-    pub memory_protection_keys: MpkEnabled,
+    pub memory_protection_keys: Enabled,
     pub max_memory_protection_keys: usize,
+
+    pub pagemap_scan: Enabled,
 }
 
 impl PoolingAllocationConfig {
@@ -66,6 +68,8 @@ impl PoolingAllocationConfig {
 
         cfg.opts.pooling_memory_protection_keys = Some(self.memory_protection_keys);
         cfg.opts.pooling_max_memory_protection_keys = Some(self.max_memory_protection_keys);
+
+        cfg.opts.pooling_pagemap_scan = Some(self.pagemap_scan);
     }
 }
 
@@ -108,8 +112,10 @@ impl<'a> Arbitrary<'a> for PoolingAllocationConfig {
 
             async_stack_keep_resident: u.int_in_range(0..=1 << 20)?,
 
-            memory_protection_keys: *u.choose(&[MpkEnabled::Auto, MpkEnabled::Disable])?,
+            memory_protection_keys: *u.choose(&[Enabled::Auto, Enabled::No])?,
             max_memory_protection_keys: u.int_in_range(1..=20)?,
+
+            pagemap_scan: *u.choose(&[Enabled::Auto, Enabled::No])?,
         })
     }
 }

crates/wasmtime/src/config.rs

@@ -2950,16 +2950,17 @@ pub enum WasmBacktraceDetails {
     Environment,
 }
 
-/// Describe the tri-state configuration of memory protection keys (MPK).
+/// Describe the tri-state configuration of keys such as MPK or PAGEMAP_SCAN.
 #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
-pub enum MpkEnabled {
-    /// Use MPK if supported by the current system; fall back to guard regions
-    /// otherwise.
+pub enum Enabled {
+    /// Enable this feature if it's detected on the host system, otherwise leave
+    /// it disabled.
     Auto,
-    /// Use MPK or fail if not supported.
-    Enable,
-    /// Do not use MPK.
-    Disable,
+    /// Enable this feature and fail configuration if the feature is not
+    /// detected on the host system.
+    Yes,
+    /// Do not enable this feature, even if the host system supports it.
+    No,
 }
 
 /// Configuration options used with [`InstanceAllocationStrategy::Pooling`] to
@@ -3482,19 +3483,19 @@ impl PoolingAllocationConfig {
     ///
     /// - `auto`: if MPK support is available the guard regions are removed; if
     ///   not, the guard regions remain
-    /// - `enable`: use MPK to eliminate guard regions; fail if MPK is not
+    /// - `yes`: use MPK to eliminate guard regions; fail if MPK is not
     ///   supported
-    /// - `disable`: never use MPK
+    /// - `no`: never use MPK
     ///
-    /// By default this value is `disabled`, but may become `auto` in future
+    /// By default this value is `no`, but may become `auto` in future
     /// releases.
     ///
     /// __WARNING__: this configuration options is still experimental--use at
     /// your own risk! MPK uses kernel and CPU features to protect memory
     /// regions; you may observe segmentation faults if anything is
     /// misconfigured.
     #[cfg(feature = "memory-protection-keys")]
-    pub fn memory_protection_keys(&mut self, enable: MpkEnabled) -> &mut Self {
+    pub fn memory_protection_keys(&mut self, enable: Enabled) -> &mut Self {
         self.config.memory_protection_keys = enable;
         self
     }
@@ -3520,7 +3521,7 @@ impl PoolingAllocationConfig {
     /// Check if memory protection keys (MPK) are available on the current host.
     ///
     /// This is a convenience method for determining MPK availability using the
-    /// same method that [`MpkEnabled::Auto`] does. See
+    /// same method that [`Enabled::Auto`] does. See
     /// [`PoolingAllocationConfig::memory_protection_keys`] for more
     /// information.
     #[cfg(feature = "memory-protection-keys")]
@@ -3541,6 +3542,45 @@ impl PoolingAllocationConfig {
         self.config.limits.total_gc_heaps = count;
         self
     }
+
+    /// Configures whether the Linux-specific [`PAGEMAP_SCAN` ioctl][ioctl] is
+    /// used to help reset linear memory.
+    ///
+    /// When [`Self::linear_memory_keep_resident`] or
+    /// [`Self::table_keep_resident`] options are configured to nonzero values
+    /// the default behavior is to `memset` the lowest addresses of a table or
+    /// memory back to their original contents. With the `PAGEMAP_SCAN` ioctl on
+    /// Linux this can be done to more intelligently scan for resident pages in
+    /// the region and only reset those pages back to their original contents
+    /// with `memset` rather than assuming the low addresses are all resident.
+    ///
+    /// This ioctl has the potential to provide a number of performance benefits
+    /// in high-reuse and high concurrency scenarios. Notably this enables
+    /// Wasmtime to scan the entire region of WebAssembly linear memory and
+    /// manually reset memory back to its original contents, up to
+    /// [`Self::linear_memory_keep_resident`] bytes, possibly skipping an
+    /// `madvise` entirely. This can be more efficient by avoiding removing
+    /// pages from the address space entirely and additionally ensuring that
+    /// future use of the linear memory doesn't incur page faults as the pages
+    /// remain resident.
+    ///
+    /// At this time this configuration option is still being evaluated as to
+    /// how appropriate it is for all use cases. It currently defaults to
+    /// `no` or disabled but may change to `auto`, enable if supported, in the
+    /// future. This option is only supported on Linux and requires a kernel
+    /// version of 6.7 or higher.
+    ///
+    /// [ioctl]: https://www.man7.org/linux/man-pages/man2/PAGEMAP_SCAN.2const.html
+    pub fn pagemap_scan(&mut self, enable: Enabled) -> &mut Self {
+        self.config.pagemap_scan = enable;
+        self
+    }
+
+    /// Tests whether [`Self::pagemap_scan`] is available or not on the host
+    /// system.
+    pub fn is_pagemap_scan_available() -> bool {
+        crate::runtime::vm::PoolingInstanceAllocatorConfig::is_pagemap_scan_available()
+    }
 }
 
 #[cfg(feature = "std")]

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

@@ -46,7 +46,7 @@ use self::table_pool::TablePool;
 use super::{
     InstanceAllocationRequest, InstanceAllocatorImpl, MemoryAllocationIndex, TableAllocationIndex,
 };
-use crate::MpkEnabled;
+use crate::Enabled;
 use crate::prelude::*;
 use crate::runtime::vm::{
     CompiledModuleId, Memory, Table,
@@ -222,9 +222,11 @@ pub struct PoolingInstanceAllocatorConfig {
     /// Same as `linear_memory_keep_resident` but for tables.
     pub table_keep_resident: usize,
     /// Whether to enable memory protection keys.
-    pub memory_protection_keys: MpkEnabled,
+    pub memory_protection_keys: Enabled,
     /// How many memory protection keys to allocate.
     pub max_memory_protection_keys: usize,
+    /// Whether to enable PAGEMAP_SCAN on Linux.
+    pub pagemap_scan: Enabled,
 }
 
 impl Default for PoolingInstanceAllocatorConfig {
@@ -239,12 +241,19 @@ impl Default for PoolingInstanceAllocatorConfig {
             async_stack_keep_resident: 0,
             linear_memory_keep_resident: 0,
             table_keep_resident: 0,
-            memory_protection_keys: MpkEnabled::Disable,
+            memory_protection_keys: Enabled::No,
             max_memory_protection_keys: 16,
+            pagemap_scan: Enabled::No,
         }
     }
 }
 
+impl PoolingInstanceAllocatorConfig {
+    pub fn is_pagemap_scan_available() -> bool {
+        PageMap::new().is_some()
+    }
+}
+
 /// An error returned when the pooling allocator cannot allocate a table,
 /// memory, etc... because the maximum number of concurrent allocations for that
 /// entity has been reached.
@@ -353,7 +362,16 @@ impl PoolingInstanceAllocator {
             gc_heaps: GcHeapPool::new(config)?,
             #[cfg(feature = "async")]
             stacks: StackPool::new(config)?,
-            pagemap: PageMap::new(),
+            pagemap: match config.pagemap_scan {
+                Enabled::Auto => PageMap::new(),
+                Enabled::Yes => Some(PageMap::new().ok_or_else(|| {
+                    anyhow!(
+                        "required to enable PAGEMAP_SCAN but this system \
+                         does not support it"
+                    )
+                })?),
+                Enabled::No => None,
+            },
         })
     }
 

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

@@ -60,7 +60,7 @@ use crate::runtime::vm::{
     MemoryImageSlot, Mmap, MmapOffset, PoolingInstanceAllocatorConfig, mmap::AlignedLength,
 };
 use crate::{
-    MpkEnabled,
+    Enabled,
     runtime::vm::mpk::{self, ProtectionKey, ProtectionMask},
     vm::HostAlignedByteCount,
 };
@@ -151,21 +151,21 @@ impl MemoryPool {
             );
         }
         let pkeys = match config.memory_protection_keys {
-            MpkEnabled::Auto => {
+            Enabled::Auto => {
                 if mpk::is_supported() {
                     mpk::keys(config.max_memory_protection_keys)
                 } else {
                     &[]
                 }
             }
-            MpkEnabled::Enable => {
+            Enabled::Yes => {
                 if mpk::is_supported() {
                     mpk::keys(config.max_memory_protection_keys)
                 } else {
                     bail!("mpk is disabled on this system")
                 }
             }
-            MpkEnabled::Disable => &[],
+            Enabled::No => &[],
         };
 
         // This is a tricky bit of global state: when creating a memory pool
@@ -840,7 +840,7 @@ mod tests {
 
         // Force the use of MPK.
         let config = PoolingInstanceAllocatorConfig {
-            memory_protection_keys: MpkEnabled::Enable,
+            memory_protection_keys: Enabled::Yes,
             ..PoolingInstanceAllocatorConfig::default()
         };
         let pool = MemoryPool::new(&config, &Tunables::default_host()).unwrap();

crates/wasmtime/tests/engine_across_forks.rs

@@ -51,6 +51,7 @@ mod unix {
     fn pooling_allocator_reset() -> Result<()> {
         let mut pooling = PoolingAllocationConfig::new();
         pooling.linear_memory_keep_resident(4096);
+        pooling.pagemap_scan(Enabled::Auto);
         let mut config = Config::new();
         config.allocation_strategy(pooling);
         config.macos_use_mach_ports(false);

examples/mpk.rs

@@ -43,11 +43,11 @@ fn main() -> Result<()> {
     let args = Args::parse();
     info!("{args:?}");
 
-    let without_mpk = probe_engine_size(&args, MpkEnabled::Disable)?;
+    let without_mpk = probe_engine_size(&args, Enabled::No)?;
     println!("without MPK:\t{}", without_mpk.to_string());
 
     if PoolingAllocationConfig::are_memory_protection_keys_available() {
-        let with_mpk = probe_engine_size(&args, MpkEnabled::Enable)?;
+        let with_mpk = probe_engine_size(&args, Enabled::Yes)?;
         println!("with MPK:\t{}", with_mpk.to_string());
         println!(
             "\t\t{}x more slots per reserved memory",
@@ -90,7 +90,7 @@ fn parse_byte_size(value: &str) -> Result<u64> {
 
 /// Find the engine with the largest number of memories we can create on this
 /// machine.
-fn probe_engine_size(args: &Args, mpk: MpkEnabled) -> Result<Pool> {
+fn probe_engine_size(args: &Args, mpk: Enabled) -> Result<Pool> {
     let mut search = ExponentialSearch::new();
     let mut mapped_bytes = 0;
     while !search.done() {
@@ -182,7 +182,7 @@ impl ExponentialSearch {
 }
 
 /// Build a pool-allocated engine with `num_memories` slots.
-fn build_engine(args: &Args, num_memories: u32, enable_mpk: MpkEnabled) -> Result<usize> {
+fn build_engine(args: &Args, num_memories: u32, enable_mpk: Enabled) -> Result<usize> {
     // Configure the memory pool.
     let mut pool = PoolingAllocationConfig::default();
     let max_memory_size =