Use MemFdSlot in the on-demand allocator as well.

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,
         }
     }