diff --git a/crates/runtime/src/mmap.rs b/crates/runtime/src/mmap.rs index 7c2fb99cd52d..24fa1351a7b0 100644 --- a/crates/runtime/src/mmap.rs +++ b/crates/runtime/src/mmap.rs @@ -424,6 +424,33 @@ impl Mmap { pub fn original_file(&self) -> Option<&Arc> { 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) -> 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 { diff --git a/crates/runtime/src/mmap_vec.rs b/crates/runtime/src/mmap_vec.rs index 2779668b20f0..a952bee7d662 100644 --- a/crates/runtime/src/mmap_vec.rs +++ b/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 { diff --git a/crates/wasmtime/src/config.rs b/crates/wasmtime/src/config.rs index ed01cca8122f..388506ba75c2 100644 --- a/crates/wasmtime/src/config.rs +++ b/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> { #[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); diff --git a/crates/wasmtime/src/module.rs b/crates/wasmtime/src/module.rs index f47fc3324b22..06b717e2a41c 100644 --- a/crates/wasmtime/src/module.rs +++ b/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