Use MemFdSlot in the on-demand allocator as well.

crates/runtime/src/instance/allocator.rs

@@ -392,15 +392,15 @@ fn initialize_memories(
     initializers: &[MemoryInitializer],
 ) -> Result<(), InstantiationError> {
     for init in initializers {
-        // Check whether this is a MemFD memory; if so, we can skip
-        // all 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].is_memfd_with_image() {
+            if !instance.memories[defined_index].needs_init() {
                 continue;
             }
         }
@@ -458,11 +458,10 @@ 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() {
+                // 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;
                 }
 
@@ -682,6 +681,7 @@ impl OnDemandInstanceAllocator {
         &self,
         module: &Module,
         store: &mut StorePtr,
+        memfds: &Option<Arc<ModuleMemFds>>,
     ) -> Result<PrimaryMap<DefinedMemoryIndex, Memory>, InstantiationError> {
         let creator = self
             .mem_creator
@@ -690,13 +690,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)?,
             );
         }
@@ -719,7 +732,7 @@ 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(()));

crates/runtime/src/lib.rs

@@ -70,6 +70,8 @@ pub use module_id::{CompiledModuleId, CompiledModuleIdAllocator};
 #[cfg(feature = "memfd-allocator")]
 mod memfd;
 
+pub use crate::memfd::MemoryMemFd;
+
 /// When memfd support is not included, provide a shim type and
 /// constructor instead so that higher-level code does not need
 /// feature-conditional compilation.

crates/runtime/src/memfd.rs

@@ -30,17 +30,20 @@ impl ModuleMemFds {
 
 /// One backing image for one memory.
 #[derive(Debug)]
-pub(crate) struct MemoryMemFd {
-    pub(crate) fd: Memfd,
+pub struct MemoryMemFd {
+    /// The actual memfd image: an anonymous file in memory which we
+    /// use as the backing content for a copy-on-write (CoW) mapping
+    /// in the memory region.
+    pub fd: Memfd,
     /// Length of image. Note that initial memory size may be larger;
     /// leading and trailing zeroes are truncated (handled by
     /// anonymous backing memfd).
-    pub(crate) len: usize,
+    pub len: usize,
     /// Image starts this many bytes into heap space. Note that the
     /// memfd's offsets are always equal to the heap offsets, so we
     /// map at an offset into the fd as well. (This simplifies
     /// construction.)
-    pub(crate) offset: usize,
+    pub offset: usize,
 }
 
 fn unsupported_initializer(segment: &MemoryInitializer, plan: &MemoryPlan) -> bool {

crates/runtime/src/memory.rs

@@ -3,6 +3,7 @@
 //! `RuntimeLinearMemory` is to WebAssembly linear memories what `Table` is to WebAssembly tables.
 
 use crate::instance::MemFdSlot;
+use crate::memfd::MemoryMemFd;
 use crate::mmap::Mmap;
 use crate::vmcontext::VMMemoryDefinition;
 use crate::Store;
@@ -10,6 +11,7 @@ use anyhow::Error;
 use anyhow::{bail, format_err, Result};
 use more_asserts::{assert_ge, assert_le};
 use std::convert::TryFrom;
+use std::sync::Arc;
 use wasmtime_environ::{MemoryPlan, MemoryStyle, WASM32_MAX_PAGES, WASM64_MAX_PAGES};
 
 const WASM_PAGE_SIZE: usize = wasmtime_environ::WASM_PAGE_SIZE as usize;
@@ -23,6 +25,8 @@ pub trait RuntimeMemoryCreator: Send + Sync {
         plan: &MemoryPlan,
         minimum: usize,
         maximum: Option<usize>,
+        // Optionally, a memfd image for CoW backing.
+        memfd_image: Option<&Arc<MemoryMemFd>>,
     ) -> Result<Box<dyn RuntimeLinearMemory>>;
 }
 
@@ -36,8 +40,14 @@ impl RuntimeMemoryCreator for DefaultMemoryCreator {
         plan: &MemoryPlan,
         minimum: usize,
         maximum: Option<usize>,
+        memfd_image: Option<&Arc<MemoryMemFd>>,
     ) -> Result<Box<dyn RuntimeLinearMemory>> {
-        Ok(Box::new(MmapMemory::new(plan, minimum, maximum)?))
+        Ok(Box::new(MmapMemory::new(
+            plan,
+            minimum,
+            maximum,
+            memfd_image,
+        )?))
     }
 }
 
@@ -59,6 +69,11 @@ pub trait RuntimeLinearMemory: Send + Sync {
     /// Return a `VMMemoryDefinition` for exposing the memory to compiled wasm
     /// code.
     fn vmmemory(&self) -> VMMemoryDefinition;
+
+    /// Does this memory need initialization? It may not if it already
+    /// has initial contents courtesy of the `MemoryMemFd` passed to
+    /// `RuntimeMemoryCreator::new_memory()`.
+    fn needs_init(&self) -> bool;
 }
 
 /// A linear memory instance.
@@ -87,11 +102,24 @@ pub struct MmapMemory {
     // optimize loads and stores with constant offsets.
     pre_guard_size: usize,
     offset_guard_size: usize,
+
+    // A MemFd CoW mapping that provides the initial content of this
+    // MmapMemory, if mapped.
+    //
+    // N.B.: this comes after the `mmap` field above because it must
+    // be destructed first. It puts a placeholder mapping in place on
+    // drop, then the `mmap` above completely unmaps the region.
+    memfd: Option<MemFdSlot>,
 }
 
 impl MmapMemory {
     /// Create a new linear memory instance with specified minimum and maximum number of wasm pages.
-    pub fn new(plan: &MemoryPlan, minimum: usize, mut maximum: Option<usize>) -> Result<Self> {
+    pub fn new(
+        plan: &MemoryPlan,
+        minimum: usize,
+        mut maximum: Option<usize>,
+        memfd_image: Option<&Arc<MemoryMemFd>>,
+    ) -> Result<Self> {
         // It's a programmer error for these two configuration values to exceed
         // the host available address space, so panic if such a configuration is
         // found (mostly an issue for hypothetical 32-bit hosts).
@@ -127,6 +155,18 @@ impl MmapMemory {
             mmap.make_accessible(pre_guard_bytes, minimum)?;
         }
 
+        // If a memfd image was specified, try to create the MemFdSlot on top of our mmap.
+        let memfd = match memfd_image {
+            Some(image) => {
+                let base = unsafe { mmap.as_mut_ptr().offset(pre_guard_bytes as isize) };
+                let len = request_bytes - pre_guard_bytes;
+                let mut memfd_slot = MemFdSlot::create(base as *mut _, len);
+                memfd_slot.instantiate(minimum, Some(image))?;
+                Some(memfd_slot)
+            }
+            None => None,
+        };
+
         Ok(Self {
             mmap,
             accessible: minimum,
@@ -134,6 +174,7 @@ impl MmapMemory {
             pre_guard_size: pre_guard_bytes,
             offset_guard_size: offset_guard_bytes,
             extra_to_reserve_on_growth,
+            memfd,
         })
     }
 }
@@ -166,7 +207,19 @@ impl RuntimeLinearMemory for MmapMemory {
             new_mmap.as_mut_slice()[self.pre_guard_size..][..self.accessible]
                 .copy_from_slice(&self.mmap.as_slice()[self.pre_guard_size..][..self.accessible]);
 
+            // Now drop the MemFdSlot, if any. We've lost the CoW
+            // advantages by explicitly copying all data, but we have
+            // preserved all of its content; so we no longer need the
+            // memfd mapping. We need to do this before we
+            // (implicitly) drop the `mmap` field by overwriting it
+            // below.
+            let _ = self.memfd.take();
+
             self.mmap = new_mmap;
+        } else if let Some(memfd) = self.memfd.as_mut() {
+            // MemFdSlot has its own growth mechanisms; defer to its
+            // implementation.
+            memfd.set_heap_limit(new_size)?;
         } else {
             // If the new size of this heap fits within the existing allocation
             // then all we need to do is to make the new pages accessible. This
@@ -192,6 +245,12 @@ impl RuntimeLinearMemory for MmapMemory {
             current_length: self.accessible,
         }
     }
+
+    fn needs_init(&self) -> bool {
+        // If we're using a memfd CoW mapping, then no initialization
+        // is needed.
+        self.memfd.is_none()
+    }
 }
 
 /// Representation of a runtime wasm linear memory.
@@ -232,9 +291,15 @@ impl Memory {
         plan: &MemoryPlan,
         creator: &dyn RuntimeMemoryCreator,
         store: &mut dyn Store,
+        memfd_image: Option<&Arc<MemoryMemFd>>,
     ) -> Result<Self> {
         let (minimum, maximum) = Self::limit_new(plan, store)?;
-        Ok(Memory::Dynamic(creator.new_memory(plan, minimum, maximum)?))
+        Ok(Memory::Dynamic(creator.new_memory(
+            plan,
+            minimum,
+            maximum,
+            memfd_image,
+        )?))
     }
 
     /// Create a new static (immovable) memory instance for the specified plan.
@@ -382,19 +447,17 @@ impl Memory {
         }
     }
 
-    /// Returns whether or not this memory is backed by a MemFD
-    /// image. Note that this is testing whether there is actually an
-    /// *image* mapped, not just whether the MemFD mechanism is being
-    /// used. The distinction is important because if we are not using
-    /// a prevalidated and prepared image, we need to fall back to
-    /// ordinary initialization code.
-    pub(crate) fn is_memfd_with_image(&self) -> bool {
+    /// Returns whether or not this memory needs initialization. It
+    /// may not if it already has initial content thanks to a CoW
+    /// mechanism like memfd.
+    pub(crate) fn needs_init(&self) -> bool {
         match self {
             Memory::Static {
                 memfd_slot: Some(ref slot),
                 ..
-            } => slot.has_image(),
-            _ => false,
+            } => !slot.has_image(),
+            Memory::Dynamic(mem) => mem.needs_init(),
+            _ => true,
         }
     }
 

crates/wasmtime/src/trampoline/memory.rs

@@ -6,7 +6,9 @@ use anyhow::{anyhow, Result};
 use std::convert::TryFrom;
 use std::sync::Arc;
 use wasmtime_environ::{EntityIndex, MemoryPlan, MemoryStyle, Module, WASM_PAGE_SIZE};
-use wasmtime_runtime::{RuntimeLinearMemory, RuntimeMemoryCreator, VMMemoryDefinition};
+use wasmtime_runtime::{
+    MemoryMemFd, RuntimeLinearMemory, RuntimeMemoryCreator, VMMemoryDefinition,
+};
 
 pub fn create_memory(store: &mut StoreOpaque, memory: &MemoryType) -> Result<InstanceId> {
     let mut module = Module::new();
@@ -46,6 +48,10 @@ impl RuntimeLinearMemory for LinearMemoryProxy {
             current_length: self.mem.byte_size(),
         }
     }
+
+    fn needs_init(&self) -> bool {
+        true
+    }
 }
 
 #[derive(Clone)]
@@ -57,6 +63,7 @@ impl RuntimeMemoryCreator for MemoryCreatorProxy {
         plan: &MemoryPlan,
         minimum: usize,
         maximum: Option<usize>,
+        _: Option<&Arc<MemoryMemFd>>,
     ) -> Result<Box<dyn RuntimeLinearMemory>> {
         let ty = MemoryType::from_wasmtime_memory(&plan.memory);
         let reserved_size_in_bytes = match plan.style {