add a config field that allows us to mlock() module mmaps into RAM

crates/runtime/src/mmap.rs

@@ -424,6 +424,33 @@ impl Mmap {
     pub fn original_file(&self) -> Option<&Arc<File>> {
         self.file.as_ref()
     }
+
+    /// Applies the `MLOCK_ONFAULT` flag to the `range` specified in this
+    /// mapping.
+    ///
+    /// This method, while available on all platforms, will unconditionally
+    /// return an error on non-Linux platforms since `MLOCK_ONFAULT` is
+    /// Linux-specific.
+    pub fn linux_mlock_onfault(&self, range: &Range<usize>) -> Result<()> {
+        assert!(range.start <= self.len());
+        assert!(range.end <= self.len());
+        assert!(range.start <= range.end);
+        assert!(
+            range.start % region::page::size() == 0,
+            "changing of protections isn't page-aligned",
+        );
+        #[cfg(target_os = "linux")]
+        unsafe {
+            rustix::io::mlock_with(
+                self.as_ptr().add(range.start) as *mut _,
+                range.end - range.start,
+                rustix::io::MlockFlags::ONFAULT,
+            )?;
+            Ok(())
+        }
+        #[cfg(not(target_os = "linux"))]
+        anyhow::bail!("mlock on fault not supported on this platform");
+    }
 }
 
 impl Drop for Mmap {

crates/runtime/src/mmap_vec.rs

@@ -128,6 +128,13 @@ impl MmapVec {
     pub fn original_offset(&self) -> usize {
         self.range.start
     }
+
+    /// Applies `MLOCK_ONFAULT` for this entire mapping.
+    ///
+    /// See [`Mmap::linux_mlock_onfault`] for more information.
+    pub fn linux_mlock_onfault(&self) -> Result<()> {
+        self.mmap.linux_mlock_onfault(&self.range)
+    }
 }
 
 impl Deref for MmapVec {

crates/wasmtime/src/config.rs

@@ -100,6 +100,7 @@ pub struct Config {
     pub(crate) paged_memory_initialization: bool,
     pub(crate) memory_init_cow: bool,
     pub(crate) memory_guaranteed_dense_image_size: u64,
+    pub(crate) linux_mlock_modules_onfault: bool,
 }
 
 impl Config {
@@ -135,6 +136,7 @@ impl Config {
             paged_memory_initialization: cfg!(all(target_os = "linux", feature = "uffd")),
             memory_init_cow: true,
             memory_guaranteed_dense_image_size: 16 << 20,
+            linux_mlock_modules_onfault: false,
         };
         #[cfg(compiler)]
         {
@@ -1250,6 +1252,36 @@ impl Config {
         self
     }
 
+    /// Enables loaded modules to be `mlock`'d into memory with the
+    /// `MLOCK_ONFAULT` flag.
+    ///
+    /// This configuration option is intended to indicate to the kernel that
+    /// once a module has been paged in (such as its text section) then it must
+    /// remain in memory, never being paged out. When modules are loaded from
+    /// precompiled images on disk then the kernel might otherwise discard the
+    /// pages for the text section or other parts of the module under memory
+    /// pressure. This means that repeated executions of a module might
+    /// sometimes take longer as the contents of the text section need to be
+    /// paged back in. By using this configuration option after the first fault
+    /// happens then the module will stick around in memory forevermore.
+    ///
+    /// Note that enabling this feature also affects how memory initialization
+    /// with copy-on-write work as well (the [`Config::memory_init_cow`]
+    /// configuration option). When loading a precompiled module from disk
+    /// typically the initialization image for memory can be reused from the
+    /// file on disk, but when this configuration option is enabled then a fresh
+    /// file descriptor from `memfd_create` is always used as the memory
+    /// initialization image to force contents to always be in memory no matter
+    /// what.
+    ///
+    /// This feature is a Linux-specific configuration option. Enabling this on
+    /// non-Linux platforms will prevent all [`Module`](crate::Module)s from
+    /// being created.
+    pub fn linux_mlock_modules_onfault(&mut self, enable: bool) -> &mut Self {
+        self.linux_mlock_modules_onfault = enable;
+        self
+    }
+
     pub(crate) fn build_allocator(&self) -> Result<Box<dyn InstanceAllocator>> {
         #[cfg(feature = "async")]
         let stack_size = self.async_stack_size;
@@ -1330,6 +1362,7 @@ impl Clone for Config {
             paged_memory_initialization: self.paged_memory_initialization,
             memory_init_cow: self.memory_init_cow,
             memory_guaranteed_dense_image_size: self.memory_guaranteed_dense_image_size,
+            linux_mlock_modules_onfault: self.linux_mlock_modules_onfault,
         }
     }
 }
@@ -1362,7 +1395,16 @@ impl fmt::Debug for Config {
                 "guard_before_linear_memory",
                 &self.tunables.guard_before_linear_memory,
             )
-            .field("parallel_compilation", &self.parallel_compilation);
+            .field("parallel_compilation", &self.parallel_compilation)
+            .field(
+                "paged_memory_initialization",
+                &self.paged_memory_initialization,
+            )
+            .field("memory_init_cow", &self.memory_init_cow)
+            .field(
+                "linux_mlock_modules_onfault",
+                &self.linux_mlock_modules_onfault,
+            );
         #[cfg(compiler)]
         {
             f.field("compiler", &self.compiler);

crates/wasmtime/src/module.rs

@@ -539,6 +539,14 @@ impl Module {
                 .flat_map(|m| m.trampolines().map(|(idx, f, _)| (idx, f))),
         ));
 
+        if engine.config().linux_mlock_modules_onfault {
+            for m in modules.iter() {
+                m.mmap()
+                    .linux_mlock_onfault()
+                    .context("failed to mlock module's memory")?;
+            }
+        }
+
         let module = modules.remove(main_module);
 
         let module_upvars = module_upvars
@@ -1146,6 +1154,18 @@ fn memory_images(engine: &Engine, module: &CompiledModule) -> Result<Option<Modu
         return Ok(None);
     }
 
-    // ... otherwise logic is delegated to the `ModuleMemoryImages::new` constructor
-    ModuleMemoryImages::new(module.module(), module.wasm_data(), Some(module.mmap()))
+    // ... otherwise logic is delegated to the `ModuleMemoryImages::new`
+    // constructor. Note that we explicitly pass `None` as the backing mmap if
+    // mlocking is enabled. Enabling mlocking means that we're trying to avoid
+    // page faults hitting disk, and if the module's mmap originally came from
+    // disk then it means that the mmap for the initialization image might miss
+    // in the kernel's page cache and hit the disk anyway. To avoid this we
+    // force a memfd to be used which means that everything should stay in
+    // memory.
+    let mmap = if engine.config().linux_mlock_modules_onfault {
+        None
+    } else {
+        Some(module.mmap())
+    };
+    ModuleMemoryImages::new(module.module(), module.wasm_data(), mmap)
 }