Merge pull request #3697 from cfallin/memfd-cow

memfd/madvise-based CoW pooling allocator

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)?;
         }
     }
 
@@ -646,6 +669,7 @@ impl OnDemandInstanceAllocator {
         &self,
         module: &Module,
         store: &mut StorePtr,
+        memfds: &Option<Arc<ModuleMemFds>>,
     ) -> Result<PrimaryMap<DefinedMemoryIndex, Memory>, InstantiationError> {
         let creator = self
             .mem_creator
@@ -654,13 +678,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)?,
             );
         }
@@ -683,25 +720,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() {