Restructure signal-based dirty page tracker to be part of SharedMemory

Signed-off-by: Ludvig Liljenberg <4257730+ludfjig@users.noreply.github.com>

Author: ludfjig

src/hyperlight_host/src/hypervisor/mod.rs

@@ -512,6 +512,7 @@ pub(crate) mod tests {
     #[cfg(gdb)]
     use crate::hypervisor::DbgMemAccessHandlerCaller;
     use crate::mem::ptr::RawPtr;
+    use crate::mem::shared_mem::SharedMemory;
     use crate::sandbox::uninitialized::GuestBinary;
     #[cfg(any(crashdump, gdb))]
     use crate::sandbox::uninitialized::SandboxRuntimeConfig;
@@ -581,6 +582,10 @@ pub(crate) mod tests {
             GuestPtr::try_from(Offset::from(0))
         }?;
 
+        // We need to stop tracking dirty pages from the host side before we start the guest
+        gshm.shared_mem
+            .with_exclusivity(|e| e.stop_tracking_dirty_pages())??;
+
         let mut vm = set_up_hypervisor_partition(
             &mut gshm,
             &config,

src/hyperlight_host/src/mem/dirty_page_tracking.rs

@@ -14,36 +14,46 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
+use std::sync::Arc;
+
 use tracing::{Span, instrument};
 
 #[cfg(target_os = "linux")]
 pub use super::linux_dirty_page_tracker::LinuxDirtyPageTracker as PlatformDirtyPageTracker;
-use super::shared_mem::SharedMemory;
+use super::shared_mem::HostMapping;
 #[cfg(target_os = "windows")]
 pub use super::windows_dirty_page_tracker::WindowsDirtyPageTracker as PlatformDirtyPageTracker;
 use crate::Result;
 
 /// Trait defining the interface for dirty page tracking implementations
 pub trait DirtyPageTracking {
-    fn get_dirty_pages(self) -> Result<Vec<usize>>;
+    #[cfg(test)]
+    fn get_dirty_pages(&self) -> Result<Vec<usize>>;
+    fn uninstall(self) -> Result<Vec<usize>>;
 }
 
 /// Cross-platform dirty page tracker that delegates to platform-specific implementations
+#[derive(Debug)]
 pub struct DirtyPageTracker {
     inner: PlatformDirtyPageTracker,
 }
 
 impl DirtyPageTracker {
     /// Create a new dirty page tracker for the given shared memory
     #[instrument(skip_all, parent = Span::current(), level = "Trace")]
-    pub fn new<T: SharedMemory>(shared_memory: &T) -> Result<Self> {
-        let inner = PlatformDirtyPageTracker::new(shared_memory)?;
+    pub fn new(mapping: Arc<HostMapping>) -> Result<Self> {
+        let inner = PlatformDirtyPageTracker::new(mapping)?;
         Ok(Self { inner })
     }
 }
 
 impl DirtyPageTracking for DirtyPageTracker {
-    fn get_dirty_pages(self) -> Result<Vec<usize>> {
+    fn uninstall(self) -> Result<Vec<usize>> {
+        self.inner.stop_tracking_and_get_dirty_pages()
+    }
+
+    #[cfg(test)]
+    fn get_dirty_pages(&self) -> Result<Vec<usize>> {
         self.inner.get_dirty_pages()
     }
 }