Merge branch 'abk/run-564-initial-memory-size-bug' into 'master'

RUN-564: Mmap enough for stable memory in mem creation

When creating a memory region for wasmtime, we were limiting its size to
the maximum 32-bit memory size, but for wasm-native stable memory this
should be increased to the maximum stable memory size. 

See merge request dfinity-lab/public/ic!11158

rs/embedders/src/wasmtime_embedder.rs

@@ -29,7 +29,7 @@ use ic_replicated_state::{
 use ic_sys::PAGE_SIZE;
 use ic_types::{
     methods::{FuncRef, WasmMethod},
-    CanisterId,
+    CanisterId, MAX_STABLE_MEMORY_IN_BYTES,
 };
 use ic_wasm_types::{BinaryEncodedWasm, WasmEngineError};
 use memory_tracker::{DirtyPageTracking, PageBitmap, SigsegvMemoryTracker};
@@ -195,12 +195,15 @@ impl WasmtimeEmbedder {
             config.wasm_memory64(true);
         }
         config
-            // maximum size in bytes where a linear memory is considered
-            // static. setting this to maximum Wasm memory size will guarantee
+            // The maximum size in bytes where a linear memory is considered
+            // static. Setting this to maximum Wasm memory size will guarantee
             // the memory is always static.
-            .static_memory_maximum_size(
-                wasmtime_environ::WASM_PAGE_SIZE as u64 * wasmtime_environ::WASM32_MAX_PAGES,
-            )
+            //
+            // If there is a change in the size of the largest memories we
+            // expect to see then the changes will likely need to be coordinated
+            // with a change in how we create the memories in the implementation
+            // of `wasmtime::MemoryCreator`.
+            .static_memory_maximum_size(MAX_STABLE_MEMORY_IN_BYTES)
             .max_wasm_stack(embedder_config.max_wasm_stack_size);
 
         config

rs/embedders/src/wasmtime_embedder/host_memory.rs

@@ -29,10 +29,6 @@ fn round_up_to_os_page_size(size: usize) -> usize {
     round_up_to_page_size(size, PAGE_SIZE)
 }
 
-fn wasm_max_mem_size_in_bytes() -> usize {
-    WASM32_MAX_PAGES as usize * WASM_PAGE_SIZE as usize
-}
-
 #[derive(Hash, PartialEq, Eq)]
 pub(crate) struct MemoryStart(pub(crate) usize);
 
@@ -65,6 +61,24 @@ unsafe impl wasmtime::MemoryCreator for WasmtimeMemoryCreator {
         reserved_size_in_bytes: Option<usize>,
         guard_size: usize,
     ) -> Result<Box<dyn wasmtime::LinearMemory>, String> {
+        // We don't use the `reserved_size_in_bytes` because the size of the
+        // memory allocation is determined based on the memory type: 64-bit
+        // memories have size at most the maximum stable memory size and 32-bit
+        // memories have size at most 4GiB. So we always allocate that amount
+        // (unless the module explicitly lists a smaller maximum).
+        //
+        // If we get a `reserved_size_in_bytes` that exceeds the max stable
+        // memory size then there has been a change in our setting of the
+        // `static_memory_maximum_size` for the `wasmtime::Config` which isn't
+        // compatible with the memory sizes we expect to see here.
+        if let Some(reserved_size) = reserved_size_in_bytes {
+            assert!(
+                reserved_size <= MAX_STABLE_MEMORY_IN_BYTES as usize,
+                "Reserved bytes for wasm memory {} exceeds the maximum expected {}",
+                reserved_size,
+                MAX_STABLE_MEMORY_IN_BYTES
+            )
+        }
         let max_pages = if ty.is_64() {
             MAX_STABLE_MEMORY_IN_BYTES / (WASM_PAGE_SIZE as u64)
         } else {
@@ -73,9 +87,7 @@ unsafe impl wasmtime::MemoryCreator for WasmtimeMemoryCreator {
         let min = std::cmp::min(ty.minimum(), max_pages) as usize;
         let max = std::cmp::min(ty.maximum().unwrap_or(max_pages), max_pages) as usize;
 
-        let mem_size = reserved_size_in_bytes.unwrap_or_else(wasm_max_mem_size_in_bytes);
-
-        let mem = MmapMemory::new(mem_size, guard_size);
+        let mem = MmapMemory::new(max * WASM_PAGE_SIZE as usize, guard_size);
 
         match self.created_memories.lock() {
             Err(err) => Err(format!("Error locking map of created memories: {:?}", err)),

rs/embedders/tests/wasmtime_embedder.rs

@@ -14,8 +14,9 @@ use ic_types::{
 mod test {
     use ic_interfaces::execution_environment::{HypervisorError, TrapCode};
     use ic_registry_subnet_type::SubnetType;
+    use ic_replicated_state::canister_state::WASM_PAGE_SIZE_IN_BYTES;
     use ic_test_utilities::wasmtime_instance::DEFAULT_NUM_INSTRUCTIONS;
-    use ic_types::{methods::WasmClosure, PrincipalId};
+    use ic_types::{methods::WasmClosure, NumBytes, PrincipalId};
 
     use super::*;
 
@@ -1563,4 +1564,45 @@ mod test {
         let err = instance.run(func_ref("test_len_both")).unwrap_err();
         assert_eq!(err, Trapped(StableMemoryOutOfBounds));
     }
+
+    /// Test that stable memory access past the normal 32-bit range (including
+    /// guard pages) works properly.
+    #[test]
+    fn stable_access_beyond_32_bit_range() {
+        fn func_ref(name: &str) -> FuncRef {
+            FuncRef::Method(WasmMethod::Update(name.to_string()))
+        }
+
+        let gb = 1024 * 1024 * 1024;
+        // We'll grow stable memory to 30 GB and then try writing to the last byte.
+        let bytes_to_access = 30 * gb;
+        let max_stable_memory_in_wasm_pages = bytes_to_access / WASM_PAGE_SIZE_IN_BYTES as u64;
+        let last_byte_of_stable_memory = bytes_to_access - 1;
+
+        let wat = format!(
+            r#"
+        (module
+            (import "ic0" "stable64_grow" (func $stable_grow (param i64) (result i64)))
+            (import "ic0" "stable64_write"
+                (func $stable64_write (param $offset i64) (param $src i64) (param $size i64)))
+            (func (export "canister_update write_to_last_page")
+                ;; Grow stable memory to the maximum size
+                (i64.eq (i64.const -1) (call $stable_grow (i64.const {max_stable_memory_in_wasm_pages})))
+                (if
+                    (then unreachable)
+                )
+
+                ;; Write to the last byte of stable memory
+                (call $stable64_write (i64.const {last_byte_of_stable_memory}) (i64.const 0) (i64.const 1))
+            )
+            (memory 2 2)
+        )"#
+        );
+
+        let mut instance = WasmtimeInstanceBuilder::new()
+            .with_wat(&wat)
+            .with_canister_memory_limit(NumBytes::from(40 * gb))
+            .build();
+        instance.run(func_ref("write_to_last_page")).unwrap();
+    }
 }

rs/test_utilities/src/wasmtime_instance.rs

@@ -1,6 +1,7 @@
 use std::convert::TryFrom;
 use std::sync::Arc;
 
+use ic_base_types::NumBytes;
 use ic_config::{flag_status::FlagStatus, subnet_config::SchedulerConfig};
 use ic_embedders::{wasm_utils::compile, wasmtime_embedder::WasmtimeInstance, WasmtimeEmbedder};
 use ic_interfaces::execution_environment::{
@@ -34,6 +35,7 @@ pub struct WasmtimeInstanceBuilder {
     subnet_type: SubnetType,
     network_topology: NetworkTopology,
     config: ic_config::embedders::Config,
+    canister_memory_limit: NumBytes,
 }
 
 impl Default for WasmtimeInstanceBuilder {
@@ -47,6 +49,7 @@ impl Default for WasmtimeInstanceBuilder {
             subnet_type: SubnetType::Application,
             network_topology: NetworkTopology::default(),
             config: ic_config::embedders::Config::default(),
+            canister_memory_limit: NumBytes::from(4 << 30), // Set to 4 GiB by default
         }
     }
 }
@@ -96,6 +99,13 @@ impl WasmtimeInstanceBuilder {
         }
     }
 
+    pub fn with_canister_memory_limit(self, canister_memory_limit: NumBytes) -> Self {
+        Self {
+            canister_memory_limit,
+            ..self
+        }
+    }
+
     pub fn try_build(
         self,
     ) -> Result<WasmtimeInstance<SystemApiImpl>, (HypervisorError, SystemApiImpl)> {
@@ -135,7 +145,7 @@ impl WasmtimeInstanceBuilder {
                     self.num_instructions,
                     self.num_instructions,
                 ),
-                canister_memory_limit: ic_types::NumBytes::from(4 << 30),
+                canister_memory_limit: self.canister_memory_limit,
                 compute_allocation: ComputeAllocation::default(),
                 subnet_type: self.subnet_type,
                 execution_mode: ExecutionMode::Replicated,