fix clean up in tests

Signed-off-by: Simon Davies <[email protected]>

Author: simongdavies

src/hyperlight_host/src/mem/linux_dirty_page_tracker.rs

@@ -26,7 +26,7 @@ use log::error;
 use crate::mem::shared_mem::{HostMapping, SharedMemory};
 use crate::{Result, new_error};
 
-// Tracker metadata stored in TLS
+// Tracker metadata stored in global lock-free storage
 struct TrackerData {
     pid: u32,
     base_addr: usize,
@@ -56,7 +56,7 @@ static HANDLER_INSTALLED: AtomicBool = AtomicBool::new(false);
 
 /// Dirty page tracker for Linux
 /// This tracks which pages have been written for a memory region once new has been called
-/// It marks pages as RO and then uses SIGSEGV to detect writes to pages, then updates the page to RW and notes the page index as dirty by writing details to TLS
+/// It marks pages as RO and then uses SIGSEGV to detect writes to pages, then updates the page to RW and notes the page index as dirty by writing details to global lock-free storage
 ///
 /// A user calls get_dirty_pages to get a list of dirty pages to get details of the pages that were written to since the tracker was created
 ///
@@ -97,7 +97,6 @@ impl LinuxDirtyPageTracker {
         // Get the current process ID
         let current_pid = std::process::id();
 
-        #[cfg(any(debug_assertions, test))]
         // Check that there is not already a tracker that includes this address range
         // within the same process (virtual addresses are only unique per process)
         for guard in get_trackers().iter() {
@@ -636,7 +635,7 @@ mod tests {
                 "Failed to create first tracker for test case: {}",
                 description
             );
-            let _tracker1 = tracker1.unwrap();
+            let tracker1 = tracker1.unwrap();
 
             // Try to create overlapping tracker - this should fail
             let test_memory2 = TestSharedMemory::new(new_addr, *new_size);
@@ -654,7 +653,7 @@ mod tests {
             println!("  ✓ Correctly rejected overlap");
 
             // Clean up by dropping the tracker
-            drop(_tracker1);
+            drop(tracker1);
             println!();
         }
 
@@ -723,7 +722,7 @@ mod tests {
                 "Failed to create first tracker for non-overlap test: {}",
                 description
             );
-            let _tracker1 = tracker1.unwrap();
+            let tracker1 = tracker1.unwrap();
 
             // Try to create non-overlapping tracker - this should succeed
             let test_memory2 = TestSharedMemory::new(new_addr, *new_size);
@@ -738,12 +737,12 @@ mod tests {
                 new_addr + new_size
             );
 
-            let _tracker2 = tracker2_result.unwrap();
+            let tracker2 = tracker2_result.unwrap();
             println!("  ✓ Correctly allowed non-overlapping ranges");
 
             // Clean up
-            drop(_tracker1);
-            drop(_tracker2);
+            drop(tracker1);
+            drop(tracker2);
             println!();
         }
 
@@ -758,8 +757,8 @@ mod tests {
         let base_addr = memory_ptr as usize;
 
         // Create two non-overlapping trackers first
-        let _tracker1 = create_test_tracker(base_addr + PAGE_SIZE * 2, PAGE_SIZE * 3).unwrap();
-        let _tracker2 = create_test_tracker(base_addr + PAGE_SIZE * 8, PAGE_SIZE * 3).unwrap();
+        let tracker1 = create_test_tracker(base_addr + PAGE_SIZE * 2, PAGE_SIZE * 3).unwrap();
+        let tracker2 = create_test_tracker(base_addr + PAGE_SIZE * 8, PAGE_SIZE * 3).unwrap();
 
         // Try to create a tracker that overlaps with tracker1
         let overlap_with_1 = create_test_tracker(base_addr + PAGE_SIZE * 3, PAGE_SIZE * 3);
@@ -786,6 +785,11 @@ mod tests {
         let no_overlap = create_test_tracker(base_addr + PAGE_SIZE * 12, PAGE_SIZE * 2);
         assert!(no_overlap.is_ok(), "Should allow non-overlapping tracker");
 
+        // Explicitly drop all trackers before freeing memory
+        drop(tracker1);
+        drop(tracker2);
+        drop(no_overlap);
+
         unsafe {
             free_aligned_memory(memory_ptr, memory_size);
         }
@@ -915,7 +919,7 @@ mod tests {
     }
 
     #[test]
-    fn test_tls_cleanup_on_drop() {
+    fn test_cleanup_on_drop() {
         let (memory_ptr, memory_size) = create_aligned_memory(PAGE_SIZE * 2);
         let addr = memory_ptr as usize;
 
@@ -932,11 +936,11 @@ mod tests {
         } // tracker is dropped here
 
         // Create a new tracker for the same memory region
-        // This should work without issues if TLS was properly cleaned up
+        // This should work without issues if data was properly cleaned up
         let new_tracker = create_test_tracker(addr, memory_size);
         assert!(
             new_tracker.is_ok(),
-            "TLS not properly cleaned up on tracker drop"
+            "Data not properly cleaned up on tracker drop"
         );
 
         unsafe {
@@ -1246,6 +1250,10 @@ mod tests {
             "Tracker 2 should store the current process ID"
         );
 
+        // Explicitly drop trackers before freeing memory
+        drop(tracker1);
+        drop(tracker2);
+
         // Clean up
         unsafe {
             free_aligned_memory(memory_ptr1, memory_size1);