Merge branch 'abk/get_memory_region-bug' into 'master'

RUN-656: Fix `PageMap::get_memory_region` bug

This fixes a small bug in `PageMap::get_memory_region` where the
reported region of zeros at the end of the `PageMap` would extend into
regions actually coming from the page delta at indices below the
requested `page_index`.

This bug has no effect on production code because the method is only
used for prefetching canister memory and we only ever prefetch indices
greater than the faulting page (so the false pages at lower indices were
ignored anyway). But it needs to be fixed for prefetching to work in
both directions. 

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

Author: adambratschikaye

rs/replicated_state/src/page_map.rs

@@ -250,7 +250,7 @@ impl std::fmt::Display for PersistenceError {
 
 /// A wrapper around the raw file descriptor to be used for memory mapping the
 /// file into the Wasm heap while executing a canister.
-#[derive(Serialize, Deserialize, Clone, Debug)]
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
 pub struct FileDescriptor {
     pub fd: RawFd,
 }
@@ -266,6 +266,7 @@ pub type FileOffset = off_t;
 /// - The page maps to the checkpoint file.
 /// - The page is in the page delta of the current `PageMap`. In this case the
 ///   range is a singleton and its contents need to be copied out.
+#[derive(Debug, PartialEq)]
 pub enum MemoryRegion<'a> {
     Zeros(Range<PageIndex>),
     BackedByFile(Range<PageIndex>, FileDescriptor),

rs/replicated_state/src/page_map/checkpoint.rs

@@ -130,7 +130,7 @@ impl Mapping {
         let num_pages = (self.mmap.len() / PAGE_SIZE) as u64;
         if page_index.get() >= num_pages {
             MemoryRegion::Zeros(Range {
-                start: PageIndex::new(num_pages),
+                start: PageIndex::new(std::cmp::max(num_pages, page_range.start.get())),
                 end: page_range.end,
             })
         } else {

rs/replicated_state/src/page_map/tests.rs

@@ -4,12 +4,12 @@ use super::{
     Buffer, FileDescriptor, PageAllocator, PageAllocatorRegistry, PageDelta, PageIndex, PageMap,
     PageMapSerialization,
 };
-use crate::page_map::TestPageAllocatorFileDescriptorImpl;
+use crate::page_map::{MemoryRegion, TestPageAllocatorFileDescriptorImpl};
 use ic_sys::PAGE_SIZE;
 use ic_types::{Height, MAX_STABLE_MEMORY_IN_BYTES};
 use nix::unistd::dup;
-use std::fs::OpenOptions;
 use std::sync::Arc;
+use std::{fs::OpenOptions, ops::Range};
 
 fn assert_equal_page_maps(page_map1: &PageMap, page_map2: &PageMap) {
     assert_eq!(page_map1.num_host_pages(), page_map2.num_host_pages());
@@ -405,3 +405,90 @@ fn calc_dirty_pages_matches_actual_change() {
         )
         .unwrap()
 }
+
+#[test]
+fn zeros_region_after_delta() {
+    let mut page_map = PageMap::new_for_testing();
+    let tmp = tempfile::Builder::new()
+        .prefix("checkpoints")
+        .tempdir()
+        .unwrap();
+    let heap_file = tmp.path().join("heap");
+    let pages = &[(PageIndex::new(1), &[1u8; PAGE_SIZE])];
+    page_map.update(pages);
+
+    page_map.persist_delta(&heap_file).unwrap();
+
+    let mut page_map = PageMap::open(
+        &heap_file,
+        Height::new(0),
+        Arc::new(TestPageAllocatorFileDescriptorImpl::new()),
+    )
+    .unwrap();
+
+    let zero_range = page_map.get_memory_region(PageIndex::new(5));
+    assert_eq!(
+        MemoryRegion::Zeros(Range {
+            start: PageIndex::new(2),
+            end: PageIndex::new(u64::MAX)
+        }),
+        zero_range
+    );
+
+    let pages = &[(PageIndex::new(3), &[1u8; PAGE_SIZE])];
+    page_map.update(pages);
+
+    let zero_range = page_map.get_memory_region(PageIndex::new(5));
+    assert_eq!(
+        MemoryRegion::Zeros(Range {
+            start: PageIndex::new(4),
+            end: PageIndex::new(u64::MAX)
+        }),
+        zero_range
+    );
+}
+
+#[test]
+fn zeros_region_within_delta() {
+    let mut page_map = PageMap::new_for_testing();
+    let tmp = tempfile::Builder::new()
+        .prefix("checkpoints")
+        .tempdir()
+        .unwrap();
+    let heap_file = tmp.path().join("heap");
+    let pages = &[(PageIndex::new(1), &[1u8; PAGE_SIZE])];
+    page_map.update(pages);
+
+    page_map.persist_delta(&heap_file).unwrap();
+
+    let mut page_map = PageMap::open(
+        &heap_file,
+        Height::new(0),
+        Arc::new(TestPageAllocatorFileDescriptorImpl::new()),
+    )
+    .unwrap();
+
+    let zero_range = page_map.get_memory_region(PageIndex::new(5));
+    assert_eq!(
+        MemoryRegion::Zeros(Range {
+            start: PageIndex::new(2),
+            end: PageIndex::new(u64::MAX)
+        }),
+        zero_range
+    );
+
+    let pages = &[
+        (PageIndex::new(3), &[1u8; PAGE_SIZE]),
+        (PageIndex::new(10), &[1u8; PAGE_SIZE]),
+    ];
+    page_map.update(pages);
+
+    let zero_range = page_map.get_memory_region(PageIndex::new(5));
+    assert_eq!(
+        MemoryRegion::Zeros(Range {
+            start: PageIndex::new(4),
+            end: PageIndex::new(10)
+        }),
+        zero_range
+    );
+}