From a1caa6549aff5f1a888d3d2da509ff553b9ec750 Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Thu, 23 Sep 2021 09:54:40 +0200
Subject: [PATCH] vmm: Add page size as a parameter for
 MemoryRangeTable::from_bitmap()

This will be helpful to support the creation of a MemoryRangeTable from
virtio-mem, as it uses 2M pages.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 virtio-devices/src/vhost_user/vu_common_ctrl.rs | 2 +-
 vm-migration/src/protocol.rs                    | 3 +--
 vmm/src/memory_manager.rs                       | 2 +-
 3 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/virtio-devices/src/vhost_user/vu_common_ctrl.rs b/virtio-devices/src/vhost_user/vu_common_ctrl.rs
index 5aa7dfdec..8b5d877bd 100644
--- a/virtio-devices/src/vhost_user/vu_common_ctrl.rs
+++ b/virtio-devices/src/vhost_user/vu_common_ctrl.rs
@@ -554,7 +554,7 @@ impl VhostUserHandle {
                 let ptr = region.as_ptr() as *const u64;
                 std::slice::from_raw_parts(ptr, len).to_vec()
             };
-            Ok(MemoryRangeTable::from_bitmap(bitmap, 0))
+            Ok(MemoryRangeTable::from_bitmap(bitmap, 0, 4096))
         } else {
             Err(Error::MissingShmLogRegion)
         }
diff --git a/vm-migration/src/protocol.rs b/vm-migration/src/protocol.rs
index aaffad317..c18eeeb70 100644
--- a/vm-migration/src/protocol.rs
+++ b/vm-migration/src/protocol.rs
@@ -184,8 +184,7 @@ pub struct MemoryRangeTable {
 }
 
 impl MemoryRangeTable {
-    pub fn from_bitmap(bitmap: Vec<u64>, start_addr: u64) -> Self {
-        let page_size = 4096;
+    pub fn from_bitmap(bitmap: Vec<u64>, start_addr: u64, page_size: u64) -> Self {
         let mut table = MemoryRangeTable::default();
         let mut entry: Option<MemoryRange> = None;
         for (i, block) in bitmap.iter().enumerate() {
diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index 1e7a67455..758febc51 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -2038,7 +2038,7 @@ impl Migratable for MemoryManager {
                 .map(|(x, y)| x | y)
                 .collect();
 
-            let sub_table = MemoryRangeTable::from_bitmap(dirty_bitmap, r.gpa);
+            let sub_table = MemoryRangeTable::from_bitmap(dirty_bitmap, r.gpa, 4096);
 
             if sub_table.regions().is_empty() {
                 info!("Dirty Memory Range Table is empty");