Use bytes for maximum size of linear memory with pooling (#8628)

* Use bytes for maximum size of linear memory with pooling

This commit changes configuration of the pooling allocator to use a
byte-based unit rather than a page based unit. The previous
`PoolingAllocatorConfig::memory_pages` configuration option configures
the maximum size that a linear memory may grow to at runtime. This is an
important factor in calculation of stripes for MPK and is also a
coarse-grained knob apart from `StoreLimiter` to limit memory
consumption. This configuration option has been renamed to
`max_memory_size` and documented that it's in terms of bytes rather than
pages as before.

Additionally the documented constraint of `max_memory_size` must be
smaller than `static_memory_bound` is now additionally enforced as a
minor clean-up as part of this PR as well.

* Review comments

* Fix benchmark build

benches/instantiation.rs

@@ -219,7 +219,7 @@ fn strategies() -> impl Iterator<Item = InstanceAllocationStrategy> {
         InstanceAllocationStrategy::OnDemand,
         InstanceAllocationStrategy::Pooling({
             let mut config = PoolingAllocationConfig::default();
-            config.memory_pages(10_000);
+            config.max_memory_size(10_000 << 16);
             config
         }),
     ]

crates/cli-flags/src/lib.rs

@@ -112,6 +112,10 @@ wasmtime_option_group! {
         /// the pooling allocator.
         pub pooling_total_stacks: Option<u32>,
 
+        /// The maximum runtime size of each linear memory in the pooling
+        /// allocator, in bytes.
+        pub pooling_max_memory_size: Option<usize>,
+
         /// Whether to enable call-indirect caching.
         pub cache_call_indirects: Option<bool>,
 
@@ -605,6 +609,9 @@ impl CommonOptions {
                     if let Some(limit) = self.opts.pooling_total_stacks {
                         cfg.total_stacks(limit);
                     }
+                    if let Some(limit) = self.opts.pooling_max_memory_size {
+                        cfg.max_memory_size(limit);
+                    }
                     if let Some(enable) = self.opts.memory_protection_keys {
                         if enable {
                             cfg.memory_protection_keys(wasmtime::MpkEnabled::Enable);

crates/fuzzing/src/generators/config.rs

@@ -76,7 +76,7 @@ impl Config {
         if let InstanceAllocationStrategy::Pooling(pooling) = &mut self.wasmtime.strategy {
             // One single-page memory
             pooling.total_memories = config.max_memories as u32;
-            pooling.memory_pages = 10;
+            pooling.max_memory_size = 10 << 16;
             pooling.max_memories_per_module = config.max_memories as u32;
 
             pooling.total_tables = config.max_tables as u32;
@@ -135,7 +135,7 @@ impl Config {
             if pooling.total_memories < 1
                 || pooling.total_tables < 5
                 || pooling.table_elements < 1_000
-                || pooling.memory_pages < 900
+                || pooling.max_memory_size < (900 << 16)
                 || pooling.total_core_instances < 500
                 || pooling.core_instance_size < 64 * 1024
             {
@@ -414,20 +414,18 @@ impl<'a> Arbitrary<'a> for Config {
 
             // Ensure the pooling allocator can support the maximal size of
             // memory, picking the smaller of the two to win.
-            let min = cfg
-                .max_memory32_pages
-                .min(cfg.max_memory64_pages)
-                .min(pooling.memory_pages);
-            pooling.memory_pages = min;
-            cfg.max_memory32_pages = min;
-            cfg.max_memory64_pages = min;
+            let min_pages = cfg.max_memory32_pages.min(cfg.max_memory64_pages);
+            let min = (min_pages << 16).min(pooling.max_memory_size as u64);
+            pooling.max_memory_size = min as usize;
+            cfg.max_memory32_pages = min >> 16;
+            cfg.max_memory64_pages = min >> 16;
 
             // If traps are disallowed then memories must have at least one page
             // of memory so if we still are only allowing 0 pages of memory then
             // increase that to one here.
             if cfg.disallow_traps {
-                if pooling.memory_pages == 0 {
-                    pooling.memory_pages = 1;
+                if pooling.max_memory_size == 0 {
+                    pooling.max_memory_size = 1 << 16;
                     cfg.max_memory32_pages = 1;
                     cfg.max_memory64_pages = 1;
                 }

crates/fuzzing/src/generators/pooling_config.rs

@@ -13,7 +13,7 @@ pub struct PoolingAllocationConfig {
     pub total_tables: u32,
     pub total_stacks: u32,
 
-    pub memory_pages: u64,
+    pub max_memory_size: usize,
     pub table_elements: u32,
 
     pub component_instance_size: usize,
@@ -48,7 +48,7 @@ impl PoolingAllocationConfig {
         cfg.total_tables(self.total_tables);
         cfg.total_stacks(self.total_stacks);
 
-        cfg.memory_pages(self.memory_pages);
+        cfg.max_memory_size(self.max_memory_size);
         cfg.table_elements(self.table_elements);
 
         cfg.max_component_instance_size(self.component_instance_size);
@@ -80,7 +80,7 @@ impl<'a> Arbitrary<'a> for PoolingAllocationConfig {
         const MAX_TABLES: u32 = 100;
         const MAX_MEMORIES: u32 = 100;
         const MAX_ELEMENTS: u32 = 1000;
-        const MAX_MEMORY_PAGES: u64 = 160; // 10 MiB
+        const MAX_MEMORY_SIZE: usize = 10 * (1 << 20); // 10 MiB
         const MAX_SIZE: usize = 1 << 20; // 1 MiB
         const MAX_INSTANCE_MEMORIES: u32 = 10;
         const MAX_INSTANCE_TABLES: u32 = 10;
@@ -94,7 +94,7 @@ impl<'a> Arbitrary<'a> for PoolingAllocationConfig {
             total_tables: u.int_in_range(1..=MAX_TABLES)?,
             total_stacks: u.int_in_range(1..=MAX_COUNT)?,
 
-            memory_pages: u.int_in_range(0..=MAX_MEMORY_PAGES)?,
+            max_memory_size: u.int_in_range(0..=MAX_MEMORY_SIZE)?,
             table_elements: u.int_in_range(0..=MAX_ELEMENTS)?,
 
             component_instance_size: u.int_in_range(0..=MAX_SIZE)?,

crates/wasmtime/src/config.rs

@@ -2757,27 +2757,25 @@ impl PoolingAllocationConfig {
         self
     }
 
-    /// The maximum number of Wasm pages for any linear memory defined in a
-    /// module (default is `160`).
+    /// The maximum byte size that any WebAssembly linear memory may grow to.
     ///
-    /// The default of `160` means at most 10 MiB of host memory may be
-    /// committed for each instance.
+    /// This option defaults to 10 MiB.
     ///
-    /// If a memory's minimum page limit is greater than this value, the module
-    /// will fail to instantiate.
+    /// If a memory's minimum size is greater than this value, the module will
+    /// fail to instantiate.
     ///
-    /// If a memory's maximum page limit is unbounded or greater than this
-    /// value, the maximum will be `memory_pages` for the purpose of any
-    /// `memory.grow` instruction.
+    /// If a memory's maximum size is unbounded or greater than this value, the
+    /// maximum will be `max_memory_size` for the purpose of any `memory.grow`
+    /// instruction.
     ///
     /// This value is used to control the maximum accessible space for each
     /// linear memory of a core instance.
     ///
     /// The reservation size of each linear memory is controlled by the
     /// `static_memory_maximum_size` setting and this value cannot exceed the
     /// configured static memory maximum size.
-    pub fn memory_pages(&mut self, pages: u64) -> &mut Self {
-        self.config.limits.memory_pages = pages;
+    pub fn max_memory_size(&mut self, bytes: usize) -> &mut Self {
+        self.config.limits.max_memory_size = bytes;
         self
     }
 

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

@@ -137,8 +137,9 @@ pub struct InstanceLimits {
     /// Maximum number of linear memories per instance.
     pub max_memories_per_module: u32,
 
-    /// Maximum number of Wasm pages for each linear memory.
-    pub memory_pages: u64,
+    /// Maximum byte size of a linear memory, must be smaller than
+    /// `static_memory_reservation` in `Tunables`.
+    pub max_memory_size: usize,
 
     /// The total number of GC heaps in the pool, across all instances.
     #[cfg(feature = "gc")]
@@ -166,7 +167,7 @@ impl Default for InstanceLimits {
             // have 10k+ elements.
             table_elements: 20_000,
             max_memories_per_module: 1,
-            memory_pages: 160,
+            max_memory_size: 10 * (1 << 20), // 10 MiB
             #[cfg(feature = "gc")]
             total_gc_heaps: 1000,
         }
@@ -703,7 +704,7 @@ mod test {
         let config = PoolingInstanceAllocatorConfig {
             limits: InstanceLimits {
                 total_memories: 1,
-                memory_pages: 0x10001,
+                max_memory_size: 0x100010000,
                 ..Default::default()
             },
             ..PoolingInstanceAllocatorConfig::default()
@@ -712,13 +713,14 @@ mod test {
             PoolingInstanceAllocator::new(
                 &config,
                 &Tunables {
-                    static_memory_reservation: 65536,
+                    static_memory_reservation: 0x10000,
                     ..Tunables::default_host()
                 },
             )
             .map_err(|e| e.to_string())
             .expect_err("expected a failure constructing instance allocator"),
-            "module memory page limit of 65537 exceeds the maximum of 65536"
+            "maximum memory size of 0x100010000 bytes exceeds the configured \
+             static memory reservation of 0x10000 bytes"
         );
     }
 

crates/wasmtime/src/runtime/vm/instance/allocator/pooling/memory_pool.rs

@@ -142,14 +142,16 @@ pub struct MemoryPool {
 impl MemoryPool {
     /// Create a new `MemoryPool`.
     pub fn new(config: &PoolingInstanceAllocatorConfig, tunables: &Tunables) -> Result<Self> {
-        // The maximum module memory page count cannot exceed 65536 pages
-        if config.limits.memory_pages > 0x10000 {
+        if u64::try_from(config.limits.max_memory_size).unwrap()
+            > tunables.static_memory_reservation
+        {
             bail!(
-                "module memory page limit of {} exceeds the maximum of 65536",
-                config.limits.memory_pages
+                "maximum memory size of {:#x} bytes exceeds the configured \
+                 static memory reservation of {:#x} bytes",
+                config.limits.max_memory_size,
+                tunables.static_memory_reservation
             );
         }
-
         let pkeys = match config.memory_protection_keys {
             MpkEnabled::Auto => {
                 if mpk::is_supported() {
@@ -553,9 +555,6 @@ impl SlabConstraints {
         tunables: &Tunables,
         num_pkeys_available: usize,
     ) -> Result<Self> {
-        // The maximum size a memory can grow to in this pool.
-        let max_memory_bytes = limits.memory_pages * u64::from(WASM_PAGE_SIZE);
-
         // `static_memory_bound` is the configured number of Wasm pages for a
         // static memory slot (see `Config::static_memory_maximum_size`); even
         // if the memory never grows to this size (e.g., it has a lower memory
@@ -567,9 +566,7 @@ impl SlabConstraints {
         let expected_slot_bytes = tunables.static_memory_reservation;
 
         let constraints = SlabConstraints {
-            max_memory_bytes: max_memory_bytes
-                .try_into()
-                .context("max memory is too large")?,
+            max_memory_bytes: limits.max_memory_size,
             num_slots: limits
                 .total_memories
                 .try_into()
@@ -754,7 +751,6 @@ fn calculate(constraints: &SlabConstraints) -> Result<SlabLayout> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::runtime::vm::PoolingInstanceAllocator;
     use proptest::prelude::*;
 
     #[cfg(target_pointer_width = "64")]
@@ -767,13 +763,13 @@ mod tests {
                     max_tables_per_module: 0,
                     max_memories_per_module: 3,
                     table_elements: 0,
-                    memory_pages: 1,
+                    max_memory_size: WASM_PAGE_SIZE as usize,
                     ..Default::default()
                 },
                 ..Default::default()
             },
             &Tunables {
-                static_memory_reservation: 65536,
+                static_memory_reservation: WASM_PAGE_SIZE as u64,
                 static_memory_offset_guard_size: 0,
                 ..Tunables::default_host()
             },
@@ -794,28 +790,6 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn test_pooling_allocator_with_reservation_size_exceeded() {
-        let config = PoolingInstanceAllocatorConfig {
-            limits: InstanceLimits {
-                total_memories: 1,
-                memory_pages: 2,
-                ..Default::default()
-            },
-            ..PoolingInstanceAllocatorConfig::default()
-        };
-        let pool = PoolingInstanceAllocator::new(
-            &config,
-            &Tunables {
-                static_memory_reservation: 65536,
-                static_memory_offset_guard_size: 0,
-                ..Tunables::default_host()
-            },
-        )
-        .unwrap();
-        assert_eq!(pool.memories.layout.max_memory_bytes, 2 * 65536);
-    }
-
     #[test]
     #[cfg_attr(miri, ignore)]
     fn test_pooling_allocator_striping() {

crates/wasmtime/src/runtime/vm/instance/allocator/pooling/table_pool.rs

@@ -218,7 +218,7 @@ mod tests {
             limits: InstanceLimits {
                 total_tables: 7,
                 table_elements: 100,
-                memory_pages: 0,
+                max_memory_size: 0,
                 max_memories_per_module: 0,
                 ..Default::default()
             },

examples/mpk.rs

@@ -66,7 +66,7 @@ struct Args {
     /// The maximum number of bytes for each WebAssembly linear memory in the
     /// pool.
     #[arg(long, default_value = "128MiB", value_parser = parse_byte_size)]
-    memory_size: u64,
+    memory_size: usize,
 
     /// The maximum number of bytes a memory is considered static; see
     /// `Config::static_memory_maximum_size` for more details and the default
@@ -186,8 +186,7 @@ impl ExponentialSearch {
 fn build_engine(args: &Args, num_memories: u32, enable_mpk: MpkEnabled) -> Result<usize> {
     // Configure the memory pool.
     let mut pool = PoolingAllocationConfig::default();
-    let memory_pages = args.memory_size / u64::from(wasmtime_environ::WASM_PAGE_SIZE);
-    pool.memory_pages(memory_pages);
+    pool.max_memory_size(args.memory_size);
     pool.total_memories(num_memories)
         .memory_protection_keys(enable_mpk);
 

tests/all/async_functions.rs

@@ -344,13 +344,15 @@ async fn fuel_eventually_finishes() {
 #[tokio::test]
 async fn async_with_pooling_stacks() {
     let mut pool = crate::small_pool_config();
-    pool.total_stacks(1).memory_pages(1).table_elements(0);
+    pool.total_stacks(1)
+        .max_memory_size(1 << 16)
+        .table_elements(0);
     let mut config = Config::new();
     config.async_support(true);
     config.allocation_strategy(InstanceAllocationStrategy::Pooling(pool));
     config.dynamic_memory_guard_size(0);
     config.static_memory_guard_size(0);
-    config.static_memory_maximum_size(65536);
+    config.static_memory_maximum_size(1 << 16);
 
     let engine = Engine::new(&config).unwrap();
     let mut store = Store::new(&engine, ());
@@ -366,13 +368,16 @@ async fn async_with_pooling_stacks() {
 #[tokio::test]
 async fn async_host_func_with_pooling_stacks() -> Result<()> {
     let mut pooling = crate::small_pool_config();
-    pooling.total_stacks(1).memory_pages(1).table_elements(0);
+    pooling
+        .total_stacks(1)
+        .max_memory_size(1 << 16)
+        .table_elements(0);
     let mut config = Config::new();
     config.async_support(true);
     config.allocation_strategy(InstanceAllocationStrategy::Pooling(pooling));
     config.dynamic_memory_guard_size(0);
     config.static_memory_guard_size(0);
-    config.static_memory_maximum_size(65536);
+    config.static_memory_maximum_size(1 << 16);
 
     let mut store = Store::new(&Engine::new(&config)?, ());
     let mut linker = Linker::new(store.engine());
@@ -399,7 +404,7 @@ async fn async_mpk_protection() -> Result<()> {
     pooling
         .total_memories(10)
         .total_stacks(2)
-        .memory_pages(1)
+        .max_memory_size(1 << 16)
         .table_elements(0);
     let mut config = Config::new();
     config.async_support(true);

tests/all/instance.rs

@@ -42,7 +42,7 @@ fn linear_memory_limits() -> Result<()> {
     }
     test(&Engine::default())?;
     let mut pool = crate::small_pool_config();
-    pool.memory_pages(65536);
+    pool.max_memory_size(1 << 32);
     test(&Engine::new(Config::new().allocation_strategy(
         InstanceAllocationStrategy::Pooling(pool),
     ))?)?;