memfd/madvise-based CoW pooling allocator

.github/workflows/main.yml

@@ -310,7 +310,7 @@ jobs:
       env:
         RUST_BACKTRACE: 1
 
-    # Test uffd functionality on Linux
+    # Test Linux-specific functionality
     - run: |
         cargo test --features uffd -p wasmtime-runtime instance::allocator::pooling
         cargo test --features uffd -p wasmtime-cli pooling_allocator

Cargo.toml

@@ -21,7 +21,7 @@ path = "src/bin/wasmtime.rs"
 doc = false
 
 [dependencies]
-wasmtime = { path = "crates/wasmtime", version = "0.33.0", default-features = false, features = ['cache', 'cranelift'] }
+wasmtime = { path = "crates/wasmtime", version = "0.33.0", default-features = false, features = ['cache', 'cranelift', 'pooling-allocator', 'memfd'] }
 wasmtime-cache = { path = "crates/cache", version = "=0.33.0" }
 wasmtime-cranelift = { path = "crates/cranelift", version = "=0.33.0" }
 wasmtime-environ = { path = "crates/environ", version = "=0.33.0" }
@@ -95,6 +95,7 @@ vtune = ["wasmtime/vtune"]
 wasi-crypto = ["wasmtime-wasi-crypto"]
 wasi-nn = ["wasmtime-wasi-nn"]
 uffd = ["wasmtime/uffd"]
+pooling-allocator = ["wasmtime/pooling-allocator"]
 all-arch = ["wasmtime/all-arch"]
 posix-signals-on-macos = ["wasmtime/posix-signals-on-macos"]
 

crates/environ/src/module.rs

@@ -95,6 +95,19 @@ impl MemoryPlan {
             },
         }
     }
+
+    /// Determine whether a data segment (memory initializer) is
+    /// possibly out-of-bounds. Returns `true` if the initializer has a
+    /// dynamic location and this question cannot be resolved
+    /// pre-instantiation; hence, this method's result should not be
+    /// used to signal an error, only to exit optimized/simple fastpaths.
+    pub fn initializer_possibly_out_of_bounds(&self, init: &MemoryInitializer) -> bool {
+        match init.end() {
+            // Not statically known, so possibly out of bounds (we can't guarantee in-bounds).
+            None => true,
+            Some(end) => end > self.memory.minimum * (WASM_PAGE_SIZE as u64),
+        }
+    }
 }
 
 /// A WebAssembly linear memory initializer.
@@ -113,6 +126,16 @@ pub struct MemoryInitializer {
     pub data: Range<u32>,
 }
 
+impl MemoryInitializer {
+    /// If this initializer has a definite, static, non-overflowed end address, return it.
+    pub fn end(&self) -> Option<u64> {
+        if self.base.is_some() {
+            return None;
+        }
+        self.offset.checked_add(self.data.len() as u64)
+    }
+}
+
 /// The type of WebAssembly linear memory initialization to use for a module.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub enum MemoryInitialization {

crates/jit/src/instantiate.rs

@@ -19,7 +19,10 @@ use wasmtime_environ::{
     StackMapInformation, Trampoline, Tunables, WasmFuncType, ELF_WASMTIME_ADDRMAP,
     ELF_WASMTIME_TRAPS,
 };
-use wasmtime_runtime::{GdbJitImageRegistration, InstantiationError, VMFunctionBody, VMTrampoline};
+use wasmtime_runtime::{
+    CompiledModuleId, CompiledModuleIdAllocator, GdbJitImageRegistration, InstantiationError,
+    VMFunctionBody, VMTrampoline,
+};
 
 /// This is the name of the section in the final ELF image which contains
 /// concatenated data segments from the original wasm module.
@@ -248,6 +251,8 @@ pub struct CompiledModule {
     code: Range<usize>,
     code_memory: CodeMemory,
     dbg_jit_registration: Option<GdbJitImageRegistration>,
+    /// A unique ID used to register this module with the engine.
+    unique_id: CompiledModuleId,
 }
 
 impl CompiledModule {
@@ -271,6 +276,7 @@ impl CompiledModule {
         mmap: MmapVec,
         info: Option<CompiledModuleInfo>,
         profiler: &dyn ProfilingAgent,
+        id_allocator: &CompiledModuleIdAllocator,
     ) -> Result<Arc<Self>> {
         // Transfer ownership of `obj` to a `CodeMemory` object which will
         // manage permissions, such as the executable bit. Once it's located
@@ -312,6 +318,7 @@ impl CompiledModule {
             dbg_jit_registration: None,
             code_memory,
             meta: info.meta,
+            unique_id: id_allocator.alloc(),
         };
         ret.register_debug_and_profiling(profiler)?;
 
@@ -333,6 +340,12 @@ impl CompiledModule {
         Ok(())
     }
 
+    /// Get this module's unique ID. It is unique with respect to a
+    /// single allocator (which is ordinarily held on a Wasm engine).
+    pub fn unique_id(&self) -> CompiledModuleId {
+        self.unique_id
+    }
+
     /// Returns the underlying memory which contains the compiled module's
     /// image.
     pub fn mmap(&self) -> &MmapVec {

crates/runtime/Cargo.toml

@@ -25,6 +25,7 @@ backtrace = "0.3.61"
 lazy_static = "1.3.0"
 rand = "0.8.3"
 anyhow = "1.0.38"
+memfd = { version = "0.4.1", optional = true }
 
 [target.'cfg(target_os = "macos")'.dependencies]
 mach = "0.3.2"

crates/runtime/build.rs

@@ -10,4 +10,14 @@ fn main() {
         )
         .file("src/helpers.c")
         .compile("wasmtime-helpers");
+
+    // Check to see if we are on Linux and the `memfd` feature is
+    // active. If so, enable the `memfd` rustc cfg so `#[cfg(memfd)]`
+    // will work.
+    let os = env::var("CARGO_CFG_TARGET_OS").unwrap();
+    let is_memfd = env::var("CARGO_FEATURE_MEMFD").is_ok();
+    let is_uffd = env::var("CARGO_FEATURE_UFFD").is_ok();
+    if &os == "linux" && is_memfd && !is_uffd {
+        println!("cargo:rustc-cfg=memfd");
+    }
 }

crates/runtime/src/instance.rs

@@ -97,6 +97,29 @@ pub(crate) struct Instance {
 
 #[allow(clippy::cast_ptr_alignment)]
 impl Instance {
+    /// Helper for allocators; not a public API.
+    pub(crate) fn create_raw(
+        module: &Arc<Module>,
+        wasm_data: &'static [u8],
+        memories: PrimaryMap<DefinedMemoryIndex, Memory>,
+        tables: PrimaryMap<DefinedTableIndex, Table>,
+        host_state: Box<dyn Any + Send + Sync>,
+    ) -> Instance {
+        Instance {
+            module: module.clone(),
+            offsets: VMOffsets::new(HostPtr, &module),
+            memories,
+            tables,
+            dropped_elements: EntitySet::with_capacity(module.passive_elements.len()),
+            dropped_data: EntitySet::with_capacity(module.passive_data_map.len()),
+            host_state,
+            wasm_data,
+            vmctx: VMContext {
+                _marker: std::marker::PhantomPinned,
+            },
+        }
+    }
+
     /// Helper function to access various locations offset from our `*mut
     /// VMContext` object.
     unsafe fn vmctx_plus_offset<T>(&self, offset: u32) -> *mut T {

crates/runtime/src/instance/allocator.rs

@@ -4,23 +4,22 @@ use crate::memory::{DefaultMemoryCreator, Memory};
 use crate::table::Table;
 use crate::traphandlers::Trap;
 use crate::vmcontext::{
-    VMBuiltinFunctionsArray, VMCallerCheckedAnyfunc, VMContext, VMGlobalDefinition,
-    VMSharedSignatureIndex,
+    VMBuiltinFunctionsArray, VMCallerCheckedAnyfunc, VMGlobalDefinition, VMSharedSignatureIndex,
 };
+use crate::ModuleMemFds;
 use crate::Store;
 use anyhow::Result;
 use std::alloc;
 use std::any::Any;
 use std::convert::TryFrom;
-use std::marker;
 use std::ptr::{self, NonNull};
 use std::slice;
 use std::sync::Arc;
 use thiserror::Error;
 use wasmtime_environ::{
-    DefinedFuncIndex, DefinedMemoryIndex, DefinedTableIndex, EntityRef, EntitySet, FunctionInfo,
-    GlobalInit, HostPtr, MemoryInitialization, MemoryInitializer, Module, ModuleType, PrimaryMap,
-    SignatureIndex, TableInitializer, TrapCode, VMOffsets, WasmType, WASM_PAGE_SIZE,
+    DefinedFuncIndex, DefinedMemoryIndex, DefinedTableIndex, EntityRef, FunctionInfo, GlobalInit,
+    MemoryInitialization, MemoryInitializer, Module, ModuleType, PrimaryMap, SignatureIndex,
+    TableInitializer, TrapCode, WasmType, WASM_PAGE_SIZE,
 };
 
 #[cfg(feature = "pooling-allocator")]
@@ -39,6 +38,9 @@ pub struct InstanceAllocationRequest<'a> {
     /// The base address of where JIT functions are located.
     pub image_base: usize,
 
+    /// If using MemFD-based memories, the backing MemFDs.
+    pub memfds: Option<Arc<ModuleMemFds>>,
+
     /// Descriptors about each compiled function, such as the offset from
     /// `image_base`.
     pub functions: &'a PrimaryMap<DefinedFuncIndex, FunctionInfo>,
@@ -376,9 +378,23 @@ fn check_memory_init_bounds(
 
 fn initialize_memories(
     instance: &mut Instance,
+    module: &Module,
     initializers: &[MemoryInitializer],
 ) -> Result<(), InstantiationError> {
     for init in initializers {
+        // Check whether we can skip all initializers (due to, e.g.,
+        // memfd).
+        let memory = init.memory_index;
+        if let Some(defined_index) = module.defined_memory_index(memory) {
+            // We can only skip if there is actually a MemFD image. In
+            // some situations the MemFD image creation code will bail
+            // (e.g. due to an out of bounds data segment) and so we
+            // need to fall back on the usual initialization below.
+            if !instance.memories[defined_index].needs_init() {
+                continue;
+            }
+        }
+
         instance
             .memory_init_segment(
                 init.memory_index,
@@ -432,6 +448,13 @@ fn initialize_instance(
     match &module.memory_initialization {
         MemoryInitialization::Paged { map, out_of_bounds } => {
             for (index, pages) in map {
+                // Check whether the memory actually needs
+                // initialization. It may not if we're using a CoW
+                // mechanism like memfd.
+                if !instance.memories[index].needs_init() {
+                    continue;
+                }
+
                 let memory = instance.memory(index);
                 let slice =
                     unsafe { slice::from_raw_parts_mut(memory.base, memory.current_length) };
@@ -453,7 +476,7 @@ fn initialize_instance(
             }
         }
         MemoryInitialization::Segmented(initializers) => {
-            initialize_memories(instance, initializers)?;
+            initialize_memories(instance, module, initializers)?;
         }
     }
 
@@ -648,6 +671,7 @@ impl OnDemandInstanceAllocator {
         &self,
         module: &Module,
         store: &mut StorePtr,
+        memfds: &Option<Arc<ModuleMemFds>>,
     ) -> Result<PrimaryMap<DefinedMemoryIndex, Memory>, InstantiationError> {
         let creator = self
             .mem_creator
@@ -656,13 +680,26 @@ impl OnDemandInstanceAllocator {
         let num_imports = module.num_imported_memories;
         let mut memories: PrimaryMap<DefinedMemoryIndex, _> =
             PrimaryMap::with_capacity(module.memory_plans.len() - num_imports);
-        for plan in &module.memory_plans.values().as_slice()[num_imports..] {
+        for (memory_idx, plan) in module.memory_plans.iter().skip(num_imports) {
+            // Create a MemFdSlot if there is an image for this memory.
+            let defined_memory_idx = module
+                .defined_memory_index(memory_idx)
+                .expect("Skipped imports, should never be None");
+            let memfd_image = memfds
+                .as_ref()
+                .and_then(|memfds| memfds.get_memory_image(defined_memory_idx));
+
             memories.push(
-                Memory::new_dynamic(plan, creator, unsafe {
-                    store
-                        .get()
-                        .expect("if module has memory plans, store is not empty")
-                })
+                Memory::new_dynamic(
+                    plan,
+                    creator,
+                    unsafe {
+                        store
+                            .get()
+                            .expect("if module has memory plans, store is not empty")
+                    },
+                    memfd_image,
+                )
                 .map_err(InstantiationError::Resource)?,
             );
         }
@@ -685,25 +722,14 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
         &self,
         mut req: InstanceAllocationRequest,
     ) -> Result<InstanceHandle, InstantiationError> {
-        let memories = self.create_memories(&req.module, &mut req.store)?;
+        let memories = self.create_memories(&req.module, &mut req.store, &req.memfds)?;
         let tables = Self::create_tables(&req.module, &mut req.store)?;
 
         let host_state = std::mem::replace(&mut req.host_state, Box::new(()));
 
         let mut handle = {
-            let instance = Instance {
-                module: req.module.clone(),
-                offsets: VMOffsets::new(HostPtr, &req.module),
-                memories,
-                tables,
-                dropped_elements: EntitySet::with_capacity(req.module.passive_elements.len()),
-                dropped_data: EntitySet::with_capacity(req.module.passive_data_map.len()),
-                host_state,
-                wasm_data: &*req.wasm_data,
-                vmctx: VMContext {
-                    _marker: marker::PhantomPinned,
-                },
-            };
+            let instance =
+                Instance::create_raw(&req.module, &*req.wasm_data, memories, tables, host_state);
             let layout = instance.alloc_layout();
             let instance_ptr = alloc::alloc(layout) as *mut Instance;
             if instance_ptr.is_null() {