Snapshot only contains pages dirties since last snapshot was taken

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

Author: ludfjig

src/hyperlight_host/src/mem/mgr.rs

@@ -15,6 +15,7 @@ limitations under the License.
 */
 
 use std::cmp::Ordering;
+use std::sync::Arc;
 
 use hyperlight_common::flatbuffer_wrappers::function_call::{
     FunctionCall, validate_guest_function_call_buffer,
@@ -74,6 +75,8 @@ pub(crate) struct SandboxMemoryManager<S> {
     pub(crate) entrypoint_offset: Offset,
     /// How many memory regions were mapped after sandbox creation
     pub(crate) mapped_rgns: u64,
+    /// Most recent snapshot taken, in other words, the most recent state that `self` has been in (disregarding currently dirty pages)
+    pub(crate) most_recent_snapshot: Option<Arc<SharedMemorySnapshot>>,
 }
 
 impl<S> SandboxMemoryManager<S>
@@ -94,6 +97,7 @@ where
             load_addr,
             entrypoint_offset,
             mapped_rgns: 0,
+            most_recent_snapshot: None,
         }
     }
 
@@ -265,20 +269,32 @@ where
         &mut self,
         sandbox_id: u64,
         mapped_regions: Vec<MemoryRegion>,
-    ) -> Result<SharedMemorySnapshot> {
-        SharedMemorySnapshot::new(&mut self.shared_mem, sandbox_id, mapped_regions)
+        dirty_pages_bitmap: &[u64],
+    ) -> Result<Arc<SharedMemorySnapshot>> {
+        let snapshot = Arc::new(SharedMemorySnapshot::new(
+            &mut self.shared_mem,
+            sandbox_id,
+            mapped_regions,
+            dirty_pages_bitmap,
+            self.most_recent_snapshot.clone(),
+        )?);
+        self.most_recent_snapshot = Some(snapshot.clone());
+        Ok(snapshot)
     }
 
     /// This function restores a memory snapshot from a given snapshot.
-    pub(crate) fn restore_snapshot(&mut self, snapshot: &SharedMemorySnapshot) -> Result<()> {
-        if self.shared_mem.mem_size() != snapshot.mem_size() {
-            return Err(new_error!(
-                "Snapshot size does not match current memory size: {} != {}",
-                self.shared_mem.raw_mem_size(),
-                snapshot.mem_size()
-            ));
-        }
-        snapshot.restore_from_snapshot(&mut self.shared_mem)?;
+    pub(crate) fn restore_snapshot(
+        &mut self,
+        snapshot: &Arc<SharedMemorySnapshot>,
+        dirty_pages_bitmap: &[u64],
+    ) -> Result<()> {
+        snapshot.restore_from_snapshot(
+            &mut self.shared_mem,
+            dirty_pages_bitmap,
+            &self.most_recent_snapshot,
+        )?;
+        // Update the most recent snapshot to the one we just restored to
+        self.most_recent_snapshot = Some(snapshot.clone());
         Ok(())
     }
 
@@ -415,13 +431,15 @@ impl SandboxMemoryManager<ExclusiveSharedMemory> {
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
                 mapped_rgns: 0,
+                most_recent_snapshot: self.most_recent_snapshot.clone(),
             },
             SandboxMemoryManager {
                 shared_mem: gshm,
                 layout: self.layout,
                 load_addr: self.load_addr.clone(),
                 entrypoint_offset: self.entrypoint_offset,
                 mapped_rgns: 0,
+                most_recent_snapshot: self.most_recent_snapshot.clone(),
             },
         )
     }

src/hyperlight_host/src/mem/shared_mem.rs

@@ -634,6 +634,7 @@ impl ExclusiveSharedMemory {
 
     /// Copy the entire contents of `self` into a `Vec<u8>`, then return it
     #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    #[cfg(test)]
     pub(crate) fn copy_all_to_vec(&self) -> Result<Vec<u8>> {
         let data = self.as_slice();
         Ok(data.to_vec())