Remove dirty_pages from HostMapping struct, remove some unnecessary Options in parameters.

Signed-off-by: Ludvig Liljenberg <[email protected]>

Author: ludfjig

src/hyperlight_host/src/hypervisor/mod.rs

@@ -502,6 +502,7 @@ pub(crate) mod tests {
     use super::handlers::{MemAccessHandler, OutBHandler};
     #[cfg(gdb)]
     use crate::hypervisor::DbgMemAccessHandlerCaller;
+    use crate::mem::dirty_page_tracking::DirtyPageTracking;
     use crate::mem::ptr::RawPtr;
     use crate::sandbox::uninitialized::GuestBinary;
     #[cfg(any(crashdump, gdb))]
@@ -553,15 +554,15 @@ pub(crate) mod tests {
         let rt_cfg: SandboxRuntimeConfig = Default::default();
         let sandbox =
             UninitializedSandbox::new(GuestBinary::FilePath(filename.clone()), Some(config))?;
-        let tracker = sandbox.tracker;
+        let tracker = sandbox.tracker.unwrap();
         let (_hshm, mut gshm) = sandbox.mgr.build();
         let mut vm = set_up_hypervisor_partition(
             &mut gshm,
             &config,
             #[cfg(any(crashdump, gdb))]
             &rt_cfg,
-            tracker,
         )?;
+        let _ = tracker.get_dirty_pages()?;
         vm.initialise(
             RawPtr::from(0x230000),
             1234567890,

src/hyperlight_host/src/mem/dirty_page_tracking.rs

@@ -25,7 +25,7 @@ use crate::Result;
 
 /// Trait defining the interface for dirty page tracking implementations
 pub trait DirtyPageTracking {
-    fn get_dirty_pages(self) -> Vec<usize>;
+    fn get_dirty_pages(self) -> Result<Vec<usize>>;
 }
 
 /// Cross-platform dirty page tracker that delegates to platform-specific implementations
@@ -43,7 +43,7 @@ impl DirtyPageTracker {
 }
 
 impl DirtyPageTracking for DirtyPageTracker {
-    fn get_dirty_pages(self) -> Vec<usize> {
-        self.inner.get_dirty_pages()
+    fn get_dirty_pages(self) -> Result<Vec<usize>> {
+        self.inner.uninstall()
     }
 }

src/hyperlight_host/src/mem/linux_dirty_page_tracker.rs

@@ -163,9 +163,10 @@ impl LinuxDirtyPageTracker {
         })
     }
 
-    /// Get all dirty page indices for this tracker
-    pub(super) fn get_dirty_pages(self) -> Vec<usize> {
-        let res = if let Some(tracker_data) = get_trackers().get(&self.id) {
+    /// Get all dirty page indices for this tracker.
+    /// NOTE: This is not a bitmap, but a vector of indices where each index corresponds to a page that has been written to.
+    fn get_dirty_pages(&self) -> Result<Vec<usize>> {
+        let res: Vec<usize> = if let Some(tracker_data) = get_trackers().get(&self.id) {
             let mut dirty_pages = Vec::new();
             let tracker_data = tracker_data.val();
             for (idx, dirty) in tracker_data.dirty_pages.iter().enumerate() {
@@ -175,13 +176,20 @@ impl LinuxDirtyPageTracker {
             }
             dirty_pages
         } else {
-            Vec::new()
+            return Err(new_error!(
+                "Tried to get dirty pages from tracker, but no tracker data found"
+            ));
         };
 
-        // self is dropped here, triggering cleanup
+        Ok(res)
+    }
+
+    /// Uninstall the tracker and return the dirty pages
+    pub(super) fn uninstall(self) -> Result<Vec<usize>> {
+        let dirty_pages = self.get_dirty_pages()?;
         // explicit to document intent
         drop(self);
-        res
+        Ok(dirty_pages)
     }
 
     #[cfg(test)]
@@ -793,7 +801,7 @@ mod tests {
         let addr = memory_ptr as usize;
 
         let tracker = create_test_tracker(addr, memory_size).unwrap();
-        let dirty_pages = tracker.get_dirty_pages();
+        let dirty_pages = tracker.get_dirty_pages().unwrap();
         assert!(dirty_pages.is_empty());
 
         // tracker is already dropped by get_dirty_pages() call above
@@ -831,7 +839,7 @@ mod tests {
             }
         }
 
-        let dirty_pages = tracker.get_dirty_pages();
+        let dirty_pages = tracker.get_dirty_pages().unwrap();
 
         println!("Dirty Pages expected: {:?}", pages_to_dirty);
         println!("Dirty pages found: {:?}", dirty_pages);
@@ -876,8 +884,8 @@ mod tests {
             std::ptr::write_volatile((addr2 + PAGE_SIZE + 200) as *mut u8, 2);
         }
 
-        let dirty1 = tracker1.get_dirty_pages();
-        let dirty2 = tracker2.get_dirty_pages();
+        let dirty1 = tracker1.get_dirty_pages().unwrap();
+        let dirty2 = tracker2.get_dirty_pages().unwrap();
 
         // Verify each tracker only reports pages that were actually written to
         // Tracker1: wrote to offset 100, which is in page 0
@@ -956,7 +964,7 @@ mod tests {
             }
         }
 
-        let dirty_pages = tracker.get_dirty_pages();
+        let dirty_pages = tracker.get_dirty_pages().unwrap();
 
         // All writes to the same page should result in the same page being dirty
         if !dirty_pages.is_empty() {
@@ -1036,7 +1044,7 @@ mod tests {
                 }
 
                 // Final verification: check that ALL pages we wrote to are marked as dirty
-                let final_dirty_pages = tracker.get_dirty_pages();
+                let final_dirty_pages = tracker.get_dirty_pages().unwrap();
 
                 // Check that every page we wrote to is marked as dirty
                 for &page_idx in &pages_written {
@@ -1198,7 +1206,7 @@ mod tests {
             }
         }
 
-        let dirty_pages = tracker.get_dirty_pages();
+        let dirty_pages = tracker.get_dirty_pages().unwrap();
         println!("Stress test dirty pages: {:?}", dirty_pages);
 
         // Verify all page indices are valid

src/hyperlight_host/src/mem/mgr.rs

@@ -24,7 +24,7 @@ use hyperlight_common::flatbuffer_wrappers::function_types::ReturnValue;
 use hyperlight_common::flatbuffer_wrappers::guest_error::GuestError;
 use hyperlight_common::flatbuffer_wrappers::guest_log_data::GuestLogData;
 use hyperlight_common::flatbuffer_wrappers::host_function_details::HostFunctionDetails;
-use hyperlight_common::mem::{PAGE_SIZE_USIZE, PAGES_IN_BLOCK};
+use hyperlight_common::mem::PAGES_IN_BLOCK;
 use tracing::{Span, instrument};
 
 use super::exe::ExeInfo;
@@ -266,7 +266,8 @@ where
     /// this function will create an initial snapshot and then create the SnapshotManager
     pub(crate) fn create_initial_snapshot(
         &mut self,
-        dirty_bitmap: Option<&Vec<u64>>,
+        vm_dirty_bitmap: &[u64],
+        host_dirty_page_idx: &[usize],
         layout: &SandboxMemoryLayout,
     ) -> Result<()> {
         let mut existing_snapshot_manager = self
@@ -277,35 +278,20 @@ where
         if existing_snapshot_manager.is_some() {
             log_then_return!("Snapshot manager already initialized, not creating a new one");
         }
-        let merged_page_map = if let Some(host_dirty_page_map) = self
-            .shared_mem
-            .with_exclusivity(|e| e.get_and_clear_host_dirty_page_map())??
-        {
-            // covert vec of page indices to vec of blocks
-            // memory size is the size of the shared memory minus the 2 guard pages
-            let mut res = new_page_bitmap(self.shared_mem.mem_size() - 2 * PAGE_SIZE_USIZE, false)?;
-            for page_idx in host_dirty_page_map {
-                let block_idx = page_idx / PAGES_IN_BLOCK;
-                let bit_idx = page_idx % PAGES_IN_BLOCK;
-                res[block_idx] |= 1 << bit_idx;
-            }
 
-            // merge the host dirty page map into the dirty bitmap
-            let len = res.len();
-            let merged = bitmap_union(&res, dirty_bitmap.unwrap_or(&vec![0; len]));
-            Some(merged)
-        } else {
-            None
-        };
+        // covert vec of page indices to bitmap
+        let mut res = new_page_bitmap(self.shared_mem.raw_mem_size(), false)?;
+        for page_idx in host_dirty_page_idx {
+            let block_idx = page_idx / PAGES_IN_BLOCK;
+            let bit_idx = page_idx % PAGES_IN_BLOCK;
+            res[block_idx] |= 1 << bit_idx;
+        }
 
-        let dirty_page_map = if let Some(ref dirty_page_map) = merged_page_map {
-            Some(dirty_page_map)
-        } else {
-            dirty_bitmap
-        };
+        // merge the host dirty page map into the dirty bitmap
+        let merged = bitmap_union(&res, vm_dirty_bitmap);
 
         let snapshot_manager =
-            SharedMemorySnapshotManager::new(&mut self.shared_mem, dirty_page_map, layout)?;
+            SharedMemorySnapshotManager::new(&mut self.shared_mem, &merged, layout)?;
         existing_snapshot_manager.replace(snapshot_manager);
         Ok(())
     }
@@ -351,7 +337,7 @@ where
         }
     }
 
-    pub(crate) fn push_state(&mut self, dirty_bitmap: Option<&Vec<u64>>) -> Result<()> {
+    pub(crate) fn push_state(&mut self, dirty_bitmap: &[u64]) -> Result<()> {
         let mut snapshot_manager = self
             .snapshot_manager
             .try_lock()

src/hyperlight_host/src/mem/shared_mem.rs

@@ -19,7 +19,7 @@ use std::ffi::c_void;
 use std::io::Error;
 #[cfg(target_os = "linux")]
 use std::ptr::null_mut;
-use std::sync::{Arc, Mutex, RwLock};
+use std::sync::{Arc, RwLock};
 
 use hyperlight_common::mem::PAGE_SIZE_USIZE;
 use tracing::{Span, instrument};
@@ -96,7 +96,6 @@ pub struct HostMapping {
     size: usize,
     #[cfg(target_os = "windows")]
     handle: HANDLE,
-    dirty_pages: Mutex<Option<Vec<usize>>>,
 }
 
 impl Drop for HostMapping {
@@ -385,7 +384,6 @@ impl ExclusiveSharedMemory {
             region: Arc::new(HostMapping {
                 ptr: addr as *mut u8,
                 size: total_size,
-                dirty_pages: None.into(),
             }),
         })
     }
@@ -396,37 +394,11 @@ impl ExclusiveSharedMemory {
     }
 
     /// Stop tracking dirty pages in the shared memory region.
-    pub(crate) fn stop_tracking_dirty_pages(&self, tracker: DirtyPageTracker) -> Result<()> {
-        let dirty_pages = tracker.get_dirty_pages();
-        let mut existing_dirty_pages = self
-            .region
-            .dirty_pages
-            .lock()
-            .map_err(|e| new_error!("Failed to lock dirty_pages: {}", e))?;
-
-        match existing_dirty_pages.as_mut() {
-            Some(existing) => {
-                // merge the new dirty pages with the existing ones
-                existing.extend(dirty_pages);
-                existing.sort_unstable();
-                existing.dedup();
-            }
-            None => {
-                *existing_dirty_pages = Some(dirty_pages);
-            }
-        }
-
-        Ok(())
-    }
-
-    // Take the dirty pages
-    pub(super) fn get_and_clear_host_dirty_page_map(&self) -> Result<Option<Vec<usize>>> {
-        Ok(self
-            .region
-            .dirty_pages
-            .lock()
-            .map_err(|e| new_error!("Failed to lock dirty_pages: {}", e))?
-            .take())
+    pub(crate) fn stop_tracking_dirty_pages(
+        &self,
+        tracker: DirtyPageTracker,
+    ) -> Result<Vec<usize>> {
+        tracker.get_dirty_pages()
     }
 
     /// Create a new region of shared memory with the given minimum