memfd/madvise-based CoW pooling allocator

.github/workflows/main.yml

@@ -136,6 +136,7 @@ jobs:
     - run: cargo check -p wasmtime --no-default-features --features async
     - run: cargo check -p wasmtime --no-default-features --features uffd
     - run: cargo check -p wasmtime --no-default-features --features pooling-allocator
+    - run: cargo check -p wasmtime --no-default-features --features memfd-allocator
     - run: cargo check -p wasmtime --no-default-features --features cranelift
     - run: cargo check -p wasmtime --no-default-features --features cranelift,wat,async,cache
 
@@ -310,11 +311,13 @@ 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 test --features uffd -p wasmtime-cli wast::Cranelift
+        cargo test --features memfd-allocator -p wasmtime-cli pooling_allocator
+        cargo test --features memfd-allocator -p wasmtime-cli wast::Cranelift
       if: matrix.os == 'ubuntu-latest' && matrix.target == ''
       env:
         RUST_BACKTRACE: 1

Cargo.toml

@@ -95,6 +95,8 @@ vtune = ["wasmtime/vtune"]
 wasi-crypto = ["wasmtime-wasi-crypto"]
 wasi-nn = ["wasmtime-wasi-nn"]
 uffd = ["wasmtime/uffd"]
+pooling-allocator = ["wasmtime/pooling-allocator"]
+memfd-allocator = ["pooling-allocator", "wasmtime/memfd-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

@@ -37,6 +37,7 @@ winapi = { version = "0.3.7", features = ["winbase", "memoryapi", "errhandlingap
 
 [target.'cfg(target_os = "linux")'.dependencies]
 userfaultfd = { version = "0.4.1", optional = true }
+memfd = { version = "0.4.1", optional = true }
 
 [build-dependencies]
 cc = "1.0"
@@ -59,3 +60,5 @@ uffd = ["userfaultfd", "pooling-allocator"]
 # It is useful for applications that do not bind their own exception ports and
 # need portable signal handling.
 posix-signals-on-macos = []
+
+memfd-allocator = ["pooling-allocator", "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,28 +4,37 @@ 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")]
 mod pooling;
 
+#[cfg(feature = "memfd-allocator")]
+mod memfd;
+#[cfg(feature = "memfd-allocator")]
+pub use self::memfd::MemFdSlot;
+
+#[cfg(not(feature = "memfd-allocator"))]
+mod memfd_disabled;
+#[cfg(not(feature = "memfd-allocator"))]
+pub use self::memfd_disabled::MemFdSlot;
+
 #[cfg(feature = "pooling-allocator")]
 pub use self::pooling::{
     InstanceLimits, ModuleLimits, PoolingAllocationStrategy, PoolingInstanceAllocator,
@@ -39,6 +48,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 +388,23 @@ fn check_memory_init_bounds(
 
 fn initialize_memories(
     instance: &mut Instance,
+    module: &Module,
     initializers: &[MemoryInitializer],
 ) -> Result<(), InstantiationError> {
     for init in initializers {
+        // Check whether this is a MemFD memory; if so, we can skip
+        // all initializers.
+        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].is_memfd_with_image() {
+                continue;
+            }
+        }
+
         instance
             .memory_init_segment(
                 init.memory_index,
@@ -432,6 +458,14 @@ fn initialize_instance(
     match &module.memory_initialization {
         MemoryInitialization::Paged { map, out_of_bounds } => {
             for (index, pages) in map {
+                // 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[index].is_memfd_with_image() {
+                    continue;
+                }
+
                 let memory = instance.memory(index);
                 let slice =
                     unsafe { slice::from_raw_parts_mut(memory.base, memory.current_length) };
@@ -453,7 +487,7 @@ fn initialize_instance(
             }
         }
         MemoryInitialization::Segmented(initializers) => {
-            initialize_memories(instance, initializers)?;
+            initialize_memories(instance, module, initializers)?;
         }
     }
 
@@ -691,19 +725,8 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
         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() {