Implement a setting for reserved dynamic memory growth (#3215)

* Implement a setting for reserved dynamic memory growth

Dynamic memories aren't really that heavily used in Wasmtime right now
because for most 32-bit memories they're classified as "static" which
means they reserve 4gb of address space and never move. Growth of a
static memory is simply making pages accessible, so it's quite fast.

With the memory64 feature, however, this is no longer true since all
memory64 memories are classified as "dynamic" at this time. Previous to
this commit growth of a dynamic memory unconditionally moved the entire
linear memory in the host's address space, always resulting in a new
`Mmap` allocation. This behavior is causing fuzzers to time out when
working with 64-bit memories because incrementally growing a memory by 1
page at a time can incur a quadratic time complexity as bytes are
constantly moved.

This commit implements a scheme where there is now a tunable setting for
memory to be reserved at the end of a dynamic memory to grow into. This
means that dynamic memory growth is ideally amortized as most calls to
`memory.grow` will be able to grow into the pre-reserved space. Some
calls, though, will still need to copy the memory around.

This helps enable a commented out test for 64-bit memories now that it's
fast enough to run in debug mode. This is because the growth of memory
in the test no longer needs to copy 4gb of zeros.

* Test fixes & review comments

* More comments

crates/cranelift/src/func_environ.rs

@@ -1203,7 +1203,7 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         // allocated up front and never moved.
         let (offset_guard_size, heap_style, readonly_base) = match self.module.memory_plans[index] {
             MemoryPlan {
-                style: MemoryStyle::Dynamic,
+                style: MemoryStyle::Dynamic { .. },
                 offset_guard_size,
                 pre_guard_size: _,
                 memory: _,

crates/environ/src/module.rs

@@ -12,7 +12,10 @@ use wasmtime_types::*;
 #[derive(Debug, Clone, Hash, Serialize, Deserialize)]
 pub enum MemoryStyle {
     /// The actual memory can be resized and moved.
-    Dynamic,
+    Dynamic {
+        /// Extra space to reserve when a memory must be moved due to growth.
+        reserve: u64,
+    },
     /// Addresss space is allocated up front.
     Static {
         /// The number of mapped and unmapped pages.
@@ -54,7 +57,12 @@ impl MemoryStyle {
         }
 
         // Otherwise, make it dynamic.
-        (Self::Dynamic, tunables.dynamic_memory_offset_guard_size)
+        (
+            Self::Dynamic {
+                reserve: tunables.dynamic_memory_growth_reserve,
+            },
+            tunables.dynamic_memory_offset_guard_size,
+        )
     }
 }
 

crates/environ/src/tunables.rs

@@ -12,6 +12,12 @@ pub struct Tunables {
     /// The size in bytes of the offset guard for dynamic heaps.
     pub dynamic_memory_offset_guard_size: u64,
 
+    /// The size, in bytes, of reserved memory at the end of a "dynamic" memory,
+    /// before the guard page, that memory can grow into. This is intended to
+    /// amortize the cost of `memory.grow` in the same manner that `Vec<T>` has
+    /// space not in use to grow into.
+    pub dynamic_memory_growth_reserve: u64,
+
     /// Whether or not to generate native DWARF debug information.
     pub generate_native_debuginfo: bool,
 
@@ -66,6 +72,13 @@ impl Default for Tunables {
             // wasting too much memory.
             dynamic_memory_offset_guard_size: 0x1_0000,
 
+            // We've got lots of address space on 64-bit so use a larger
+            // grow-into-this area, but on 32-bit we aren't as lucky.
+            #[cfg(target_pointer_width = "64")]
+            dynamic_memory_growth_reserve: 2 << 30, // 2GB
+            #[cfg(target_pointer_width = "32")]
+            dynamic_memory_growth_reserve: 1 << 20, // 1MB
+
             generate_native_debuginfo: false,
             parse_wasm_debuginfo: true,
             interruptable: false,

crates/runtime/src/instance/allocator/pooling.rs

@@ -197,7 +197,7 @@ impl ModuleLimits {
                 );
             }
 
-            if let MemoryStyle::Dynamic = plan.style {
+            if let MemoryStyle::Dynamic { .. } = plan.style {
                 bail!(
                     "memory index {} has an unsupported dynamic memory plan style",
                     i,
@@ -1324,7 +1324,7 @@ mod test {
 
         let mut module = Module::default();
         module.memory_plans.push(MemoryPlan {
-            style: MemoryStyle::Dynamic,
+            style: MemoryStyle::Dynamic { reserve: 0 },
             memory: Memory {
                 minimum: 1,
                 maximum: None,

crates/runtime/src/memory.rs

@@ -77,6 +77,10 @@ pub struct MmapMemory {
     // maximum size of the linear address space reservation for this memory.
     maximum: Option<usize>,
 
+    // The amount of extra bytes to reserve whenever memory grows. This is
+    // specified so that the cost of repeated growth is amortized.
+    extra_to_reserve_on_growth: usize,
+
     // Size in bytes of extra guard pages before the start and after the end to
     // optimize loads and stores with constant offsets.
     pre_guard_size: usize,
@@ -92,15 +96,19 @@ impl MmapMemory {
         let offset_guard_bytes = usize::try_from(plan.offset_guard_size).unwrap();
         let pre_guard_bytes = usize::try_from(plan.pre_guard_size).unwrap();
 
-        let alloc_bytes = match plan.style {
+        let (alloc_bytes, extra_to_reserve_on_growth) = match plan.style {
-            MemoryStyle::Dynamic => minimum,
+            MemoryStyle::Dynamic { reserve } => (minimum, usize::try_from(reserve).unwrap()),
             MemoryStyle::Static { bound } => {
                 assert_ge!(bound, plan.memory.minimum);
-                usize::try_from(bound.checked_mul(WASM_PAGE_SIZE_U64).unwrap()).unwrap()
+                (
+                    usize::try_from(bound.checked_mul(WASM_PAGE_SIZE_U64).unwrap()).unwrap(),
+                    0,
+                )
             }
         };
         let request_bytes = pre_guard_bytes
             .checked_add(alloc_bytes)
+            .and_then(|i| i.checked_add(extra_to_reserve_on_growth))
             .and_then(|i| i.checked_add(offset_guard_bytes))
             .ok_or_else(|| format_err!("cannot allocate {} with guard regions", minimum))?;
 
@@ -115,6 +123,7 @@ impl MmapMemory {
             maximum,
             pre_guard_size: pre_guard_bytes,
             offset_guard_size: offset_guard_bytes,
+            extra_to_reserve_on_growth,
         })
     }
 }
@@ -130,11 +139,14 @@ impl RuntimeLinearMemory for MmapMemory {
 
     fn grow_to(&mut self, new_size: usize) -> Option<()> {
         if new_size > self.mmap.len() - self.offset_guard_size - self.pre_guard_size {
-            // If the new size is within the declared maximum, but needs more memory than we
+            // If the new size of this heap exceeds the current size of the
-            // have on hand, it's a dynamic heap and it can move.
+            // allocation we have, then this must be a dynamic heap. Use
+            // `new_size` to calculate a new size of an allocation, allocate it,
+            // and then copy over the memory from before.
             let request_bytes = self
                 .pre_guard_size
                 .checked_add(new_size)?
+                .checked_add(self.extra_to_reserve_on_growth)?
                 .checked_add(self.offset_guard_size)?;
 
             let mut new_mmap = Mmap::accessible_reserved(0, request_bytes).ok()?;
@@ -147,8 +159,13 @@ impl RuntimeLinearMemory for MmapMemory {
 
             self.mmap = new_mmap;
         } 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
+            // can happen either for "static" heaps which always hit this case,
+            // or "dynamic" heaps which have some space reserved after the
+            // initial allocation to grow into before the heap is moved in
+            // memory.
             assert!(new_size > self.accessible);
-            // Make the newly allocated pages accessible.
             self.mmap
                 .make_accessible(
                     self.pre_guard_size + self.accessible,

crates/wasmtime/src/config.rs

@@ -1166,6 +1166,45 @@ impl Config {
         self
     }
 
+    /// Configures the size, in bytes, of the extra virtual memory space
+    /// reserved after a "dynamic" memory for growing into.
+    ///
+    /// For the difference between static and dynamic memories, see the
+    /// [`Config::static_memory_maximum_size`]
+    ///
+    /// Dynamic memories can be relocated in the process's virtual address space
+    /// on growth and do not always reserve their entire space up-front. This
+    /// means that a growth of the memory may require movement in the address
+    /// space, which in the worst case can copy a large number of bytes from one
+    /// region to another.
+    ///
+    /// This setting configures how many bytes are reserved after the initial
+    /// reservation for a dynamic memory for growing into. A value of 0 here
+    /// means that no extra bytes are reserved and all calls to `memory.grow`
+    /// will need to relocate the wasm linear memory (copying all the bytes). A
+    /// value of 1 megabyte, however, means that `memory.grow` can allocate up
+    /// to a megabyte of extra memory before the memory needs to be moved in
+    /// linear memory.
+    ///
+    /// Note that this is a currently simple heuristic for optimizing the growth
+    /// of dynamic memories, primarily implemented for the memory64 propsal
+    /// where all memories are currently "dynamic". This is unlikely to be a
+    /// one-size-fits-all style approach and if you're an embedder running into
+    /// issues with dynamic memories and growth and are interested in having
+    /// other growth strategies available here please feel free to [open an
+    /// issue on the Wasmtime repository][issue]!
+    ///
+    /// [issue]: https://github.com/bytecodealliance/wasmtime/issues/ne
+    ///
+    /// ## Default
+    ///
+    /// For 64-bit platforms this defaults to 2GB, and for 32-bit platforms this
+    /// defaults to 1MB.
+    pub fn dynamic_memory_reserved_for_growth(&mut self, reserved: u64) -> &mut Self {
+        self.tunables.dynamic_memory_growth_reserve = round_up_to_pages(reserved);
+        self
+    }
+
     /// Indicates whether a guard region is present before allocations of
     /// linear memory.
     ///

crates/wasmtime/src/memory.rs

@@ -601,7 +601,7 @@ mod tests {
         let store = store.as_context();
         assert_eq!(store[mem.0].memory.offset_guard_size, 0);
         match &store[mem.0].memory.style {
-            wasmtime_environ::MemoryStyle::Dynamic => {}
+            wasmtime_environ::MemoryStyle::Dynamic { .. } => {}
             other => panic!("unexpected style {:?}", other),
         }
     }

crates/wasmtime/src/module/serialization.rs

@@ -429,6 +429,9 @@ impl<'a> SerializedModule<'a> {
             consume_fuel,
             static_memory_bound_is_maximum,
             guard_before_linear_memory,
+
+            // This doesn't affect compilation, it's just a runtime setting.
+            dynamic_memory_growth_reserve: _,
         } = self.tunables;
 
         Self::check_int(

crates/wasmtime/src/trampoline/memory.rs

@@ -63,7 +63,7 @@ impl RuntimeMemoryCreator for MemoryCreatorProxy {
             MemoryStyle::Static { bound } => {
                 Some(usize::try_from(bound * (WASM_PAGE_SIZE as u64)).unwrap())
             }
-            MemoryStyle::Dynamic => None,
+            MemoryStyle::Dynamic { .. } => None,
         };
         self.0
             .new_memory(

tests/all/wast.rs

@@ -69,8 +69,11 @@ fn run_wast(wast: &str, strategy: Strategy, pooling: bool) -> anyhow::Result<()>
             return Ok(());
         }
 
-        // Don't use 4gb address space reservations when not hogging memory.
+        // Don't use 4gb address space reservations when not hogging memory, and
+        // also don't reserve lots of memory after dynamic memories for growth
+        // (makes growth slower).
         cfg.static_memory_maximum_size(0);
+        cfg.dynamic_memory_reserved_for_growth(0);
     }
 
     let _pooling_lock = if pooling {

tests/misc_testsuite/memory64/more-than-4gb.wast

@@ -14,16 +14,8 @@
 )
 (assert_return (invoke "grow" (i64.const 0)) (i64.const 0x1_0001))
 (assert_return (invoke "size") (i64.const 0x1_0001))
-
-;; TODO: unsure how to test this. Right now growth of any 64-bit memory will
-;; always reallocate and copy all the previous memory to a new location, and
-;; this means that we're doing a 4gb copy here. That's pretty slow and is just
-;; copying a bunch of zeros, so until we optimize that it's not really feasible
-;; to test growth in CI andd such.
-(;
 (assert_return (invoke "grow" (i64.const 1)) (i64.const 0x1_0001))
 (assert_return (invoke "size") (i64.const 0x1_0002))
-;)
 
 ;; Test that initialization with a 64-bit global works
 (module $offset