Track dirty pages from host side

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

Author: ludfjig

src/hyperlight_host/src/mem/shared_mem.rs

@@ -19,9 +19,9 @@ use std::ffi::c_void;
 use std::io::Error;
 #[cfg(target_os = "linux")]
 use std::ptr::null_mut;
-use std::sync::{Arc, RwLock};
+use std::sync::{Arc, Mutex, RwLock};
 
-use hyperlight_common::mem::PAGE_SIZE_USIZE;
+use hyperlight_common::mem::{PAGE_SIZE_USIZE, PAGES_IN_BLOCK};
 use tracing::{Span, instrument};
 #[cfg(target_os = "windows")]
 use windows::Win32::Foundation::{CloseHandle, HANDLE, INVALID_HANDLE_VALUE};
@@ -78,8 +78,10 @@ macro_rules! generate_writer {
         #[allow(dead_code)]
         pub(crate) fn $fname(&mut self, offset: usize, value: $ty) -> Result<()> {
             let data = self.as_mut_slice();
-            bounds_check!(offset, std::mem::size_of::<$ty>(), data.len());
-            data[offset..offset + std::mem::size_of::<$ty>()].copy_from_slice(&value.to_le_bytes());
+            let size = std::mem::size_of::<$ty>();
+            bounds_check!(offset, size, data.len());
+            data[offset..offset + size].copy_from_slice(&value.to_le_bytes());
+            self.mark_pages_dirty(offset, size)?;
             Ok(())
         }
     };
@@ -133,6 +135,7 @@ impl Drop for HostMapping {
 #[derive(Debug)]
 pub struct ExclusiveSharedMemory {
     region: Arc<HostMapping>,
+    dirty_page_tracker: Arc<Mutex<Vec<u64>>>,
 }
 unsafe impl Send for ExclusiveSharedMemory {}
 
@@ -147,6 +150,8 @@ unsafe impl Send for ExclusiveSharedMemory {}
 #[derive(Debug)]
 pub struct GuestSharedMemory {
     region: Arc<HostMapping>,
+    dirty_page_tracker: Arc<Mutex<Vec<u64>>>,
+
     /// The lock that indicates this shared memory is being used by non-Rust code
     ///
     /// This lock _must_ be held whenever the guest is executing,
@@ -298,6 +303,8 @@ unsafe impl Send for GuestSharedMemory {}
 #[derive(Clone, Debug)]
 pub struct HostSharedMemory {
     region: Arc<HostMapping>,
+    dirty_page_tracker: Arc<Mutex<Vec<u64>>>,
+
     lock: Arc<RwLock<()>>,
 }
 unsafe impl Send for HostSharedMemory {}
@@ -316,6 +323,7 @@ impl ExclusiveSharedMemory {
         };
 
         use crate::error::HyperlightError::{MemoryRequestTooBig, MmapFailed, MprotectFailed};
+        use crate::mem::bitmap::new_page_bitmap;
 
         if min_size_bytes == 0 {
             return Err(new_error!("Cannot create shared memory with size 0"));
@@ -370,30 +378,48 @@ impl ExclusiveSharedMemory {
             return Err(MprotectFailed(Error::last_os_error().raw_os_error()));
         }
 
+        // HostMapping is only non-Send/Sync because raw pointers
+        // are not ("as a lint", as the Rust docs say). We don't
+        // want to mark HostMapping Send/Sync immediately, because
+        // that could socially imply that it's "safe" to use
+        // unsafe accesses from multiple threads at once. Instead, we
+        // directly impl Send and Sync on this type. Since this
+        // type does have Send and Sync manually impl'd, the Arc
+        // is not pointless as the lint suggests.
+        #[allow(clippy::arc_with_non_send_sync)]
+        let host_mapping = Arc::new(HostMapping {
+            ptr: addr as *mut u8,
+            size: total_size,
+        });
+
+        let dirty_page_tracker = new_page_bitmap(min_size_bytes, false)?;
+
         Ok(Self {
-            // HostMapping is only non-Send/Sync because raw pointers
-            // are not ("as a lint", as the Rust docs say). We don't
-            // want to mark HostMapping Send/Sync immediately, because
-            // that could socially imply that it's "safe" to use
-            // unsafe accesses from multiple threads at once. Instead, we
-            // directly impl Send and Sync on this type. Since this
-            // type does have Send and Sync manually impl'd, the Arc
-            // is not pointless as the lint suggests.
-            #[allow(clippy::arc_with_non_send_sync)]
-            region: Arc::new(HostMapping {
-                ptr: addr as *mut u8,
-                size: total_size,
-            }),
+            region: host_mapping,
+            dirty_page_tracker: Arc::new(Mutex::new(dirty_page_tracker)),
         })
     }
 
+    /// Gets the dirty bitmap and then clears it in self.
+    pub(crate) fn get_and_clear_dirty_pages(&mut self) -> Result<Vec<u64>> {
+        let mut guard = self
+            .dirty_page_tracker
+            .try_lock()
+            .map_err(|_| new_error!("Failed to acquire lock on dirty page tracker"))?;
+        let bitmap = guard.clone();
+        guard.fill(0);
+        Ok(bitmap)
+    }
+
     /// Create a new region of shared memory with the given minimum
     /// size in bytes. The region will be surrounded by guard pages.
     ///
     /// Return `Err` if shared memory could not be allocated.
     #[cfg(target_os = "windows")]
     #[instrument(skip_all, parent = Span::current(), level= "Trace")]
     pub fn new(min_size_bytes: usize) -> Result<Self> {
+        use super::bitmap::new_page_bitmap;
+
         if min_size_bytes == 0 {
             return Err(new_error!("Cannot create shared memory with size 0"));
         }
@@ -484,21 +510,26 @@ impl ExclusiveSharedMemory {
             log_then_return!(WindowsAPIError(e.clone()));
         }
 
+        // HostMapping is only non-Send/Sync because raw pointers
+        // are not ("as a lint", as the Rust docs say). We don't
+        // want to mark HostMapping Send/Sync immediately, because
+        // that could socially imply that it's "safe" to use
+        // unsafe accesses from multiple threads at once. Instead, we
+        // directly impl Send and Sync on this type. Since this
+        // type does have Send and Sync manually impl'd, the Arc
+        // is not pointless as the lint suggests.
+        #[allow(clippy::arc_with_non_send_sync)]
+        let host_mapping = Arc::new(HostMapping {
+            ptr: addr.Value as *mut u8,
+            size: total_size,
+            handle,
+        });
+
+        let dirty_page_tracker = new_page_bitmap(min_size_bytes, false)?;
+
         Ok(Self {
-            // HostMapping is only non-Send/Sync because raw pointers
-            // are not ("as a lint", as the Rust docs say). We don't
-            // want to mark HostMapping Send/Sync immediately, because
-            // that could socially imply that it's "safe" to use
-            // unsafe accesses from multiple threads at once. Instead, we
-            // directly impl Send and Sync on this type. Since this
-            // type does have Send and Sync manually impl'd, the Arc
-            // is not pointless as the lint suggests.
-            #[allow(clippy::arc_with_non_send_sync)]
-            region: Arc::new(HostMapping {
-                ptr: addr.Value as *mut u8,
-                size: total_size,
-                handle,
-            }),
+            region: host_mapping,
+            dirty_page_tracker: Arc::new(Mutex::new(dirty_page_tracker)),
         })
     }
 
@@ -613,6 +644,16 @@ impl ExclusiveSharedMemory {
         let data = self.as_mut_slice();
         bounds_check!(offset, src.len(), data.len());
         data[offset..offset + src.len()].copy_from_slice(src);
+        self.mark_pages_dirty(offset, src.len())?;
+        Ok(())
+    }
+
+    /// Copies bytes from `self` to `dst` starting at offset
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub fn copy_to_slice(&self, dst: &mut [u8], offset: usize) -> Result<()> {
+        let data = self.as_slice();
+        bounds_check!(offset, dst.len(), data.len());
+        dst.copy_from_slice(&data[offset..offset + dst.len()]);
         Ok(())
     }
 
@@ -624,6 +665,16 @@ impl ExclusiveSharedMemory {
         Ok(self.base_addr() + offset)
     }
 
+    /// Fill the memory in the range `[offset, offset + len)` with `value`
+    #[instrument(err(Debug), skip_all, parent = Span::current(), level= "Trace")]
+    pub fn zero_fill(&mut self, offset: usize, len: usize) -> Result<()> {
+        bounds_check!(offset, len, self.mem_size());
+        let data = self.as_mut_slice();
+        data[offset..offset + len].fill(0);
+        self.mark_pages_dirty(offset, len)?;
+        Ok(())
+    }
+
     generate_reader!(read_u8, u8);
     generate_reader!(read_i8, i8);
     generate_reader!(read_u16, u16);
@@ -657,15 +708,35 @@ impl ExclusiveSharedMemory {
         (
             HostSharedMemory {
                 region: self.region.clone(),
+                dirty_page_tracker: self.dirty_page_tracker.clone(),
                 lock: lock.clone(),
             },
             GuestSharedMemory {
                 region: self.region.clone(),
+                dirty_page_tracker: self.dirty_page_tracker.clone(),
                 lock: lock.clone(),
             },
         )
     }
 
+    /// Marks pages that cover bytes [offset, offset + size) as dirty
+    pub(super) fn mark_pages_dirty(&mut self, offset: usize, size: usize) -> Result<()> {
+        bounds_check!(offset, size, self.mem_size());
+        let mut bitmap = self
+            .dirty_page_tracker
+            .try_lock()
+            .map_err(|_| new_error!("Failed to lock dirty page tracker"))?;
+
+        let start_page = offset / PAGE_SIZE_USIZE;
+        let end_page = (offset + size - 1) / PAGE_SIZE_USIZE; // offset + size - 1 is the last affected byte.
+        for page_idx in start_page..=end_page {
+            let block_idx = page_idx / PAGES_IN_BLOCK;
+            let bit_idx = page_idx % PAGES_IN_BLOCK;
+            bitmap[block_idx] |= 1 << bit_idx;
+        }
+        Ok(())
+    }
+
     /// Gets the file handle of the shared memory region for this Sandbox
     #[cfg(target_os = "windows")]
     pub fn get_mmap_file_handle(&self) -> HANDLE {
@@ -743,6 +814,7 @@ impl SharedMemory for GuestSharedMemory {
     fn region(&self) -> &HostMapping {
         &self.region
     }
+
     fn with_exclusivity<T, F: FnOnce(&mut ExclusiveSharedMemory) -> T>(
         &mut self,
         f: F,
@@ -753,6 +825,7 @@ impl SharedMemory for GuestSharedMemory {
             .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
         let mut excl = ExclusiveSharedMemory {
             region: self.region.clone(),
+            dirty_page_tracker: self.dirty_page_tracker.clone(),
         };
         let ret = f(&mut excl);
         drop(excl);
@@ -803,7 +876,7 @@ impl HostSharedMemory {
     /// Write a value of type T, whose representation is the same
     /// between the sandbox and the host, and which has no invalid bit
     /// patterns
-    pub fn write<T: AllValid>(&self, offset: usize, data: T) -> Result<()> {
+    pub fn write<T: AllValid>(&mut self, offset: usize, data: T) -> Result<()> {
         bounds_check!(offset, std::mem::size_of::<T>(), self.mem_size());
         unsafe {
             let slice: &[u8] = core::slice::from_raw_parts(
@@ -812,6 +885,7 @@ impl HostSharedMemory {
             );
             self.copy_from_slice(slice, offset)?;
         }
+        self.mark_pages_dirty(offset, std::mem::size_of::<T>())?;
         Ok(())
     }
 
@@ -834,9 +908,8 @@ impl HostSharedMemory {
         Ok(())
     }
 
-    /// Copy the contents of the sandbox at the specified offset into
-    /// the slice
-    pub fn copy_from_slice(&self, slice: &[u8], offset: usize) -> Result<()> {
+    /// Copy the contents of the given slice into self
+    pub fn copy_from_slice(&mut self, slice: &[u8], offset: usize) -> Result<()> {
         bounds_check!(offset, slice.len(), self.mem_size());
         let base = self.base_ptr().wrapping_add(offset);
         let guard = self
@@ -850,6 +923,7 @@ impl HostSharedMemory {
             }
         }
         drop(guard);
+        self.mark_pages_dirty(offset, slice.len())?;
         Ok(())
     }
 
@@ -867,6 +941,7 @@ impl HostSharedMemory {
             unsafe { base.wrapping_add(i).write_volatile(value) };
         }
         drop(guard);
+        self.mark_pages_dirty(offset, len)?;
         Ok(())
     }
 
@@ -979,12 +1054,31 @@ impl HostSharedMemory {
 
         Ok(to_return)
     }
+
+    /// Marks pages that cover bytes [offset, offset + size) as dirty
+    pub(super) fn mark_pages_dirty(&mut self, offset: usize, size: usize) -> Result<()> {
+        bounds_check!(offset, size, self.mem_size());
+        let mut bitmap = self
+            .dirty_page_tracker
+            .try_lock()
+            .map_err(|_| new_error!("Failed to lock dirty page tracker"))?;
+
+        let start_page = offset / PAGE_SIZE_USIZE;
+        let end_page = (offset + size - 1) / PAGE_SIZE_USIZE; // offset + size - 1 is the last affected byte.
+        for page_idx in start_page..=end_page {
+            let block_idx = page_idx / PAGES_IN_BLOCK;
+            let bit_idx = page_idx % PAGES_IN_BLOCK;
+            bitmap[block_idx] |= 1 << bit_idx;
+        }
+        Ok(())
+    }
 }
 
 impl SharedMemory for HostSharedMemory {
     fn region(&self) -> &HostMapping {
         &self.region
     }
+
     fn with_exclusivity<T, F: FnOnce(&mut ExclusiveSharedMemory) -> T>(
         &mut self,
         f: F,
@@ -995,6 +1089,7 @@ impl SharedMemory for HostSharedMemory {
             .map_err(|e| new_error!("Error locking at {}:{}: {}", file!(), line!(), e))?;
         let mut excl = ExclusiveSharedMemory {
             region: self.region.clone(),
+            dirty_page_tracker: self.dirty_page_tracker.clone(),
         };
         let ret = f(&mut excl);
         drop(excl);
@@ -1048,7 +1143,7 @@ mod tests {
         let mem_size: usize = 4096;
         let vec_len = 10;
         let eshm = ExclusiveSharedMemory::new(mem_size)?;
-        let (hshm, _) = eshm.build();
+        let (mut hshm, _) = eshm.build();
         let vec = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
         // write the value to the memory at the beginning.
         hshm.copy_from_slice(&vec, 0)?;
@@ -1135,8 +1230,8 @@ mod tests {
     #[test]
     fn clone() {
         let eshm = ExclusiveSharedMemory::new(PAGE_SIZE_USIZE).unwrap();
-        let (hshm1, _) = eshm.build();
-        let hshm2 = hshm1.clone();
+        let (mut hshm1, _) = eshm.build();
+        let mut hshm2 = hshm1.clone();
 
         // after hshm1 is cloned, hshm1 and hshm2 should have identical
         // memory sizes and pointers.