1 files changed, 39 insertions, 18 deletions

diff --git a/crates/buddy_allocator/tests/hosted_test.rs b/crates/buddy_allocator/tests/hosted_test.rs
index a841802..d6d3389 100644
--- a/crates/buddy_allocator/tests/hosted_test.rs
+++ b/crates/buddy_allocator/tests/hosted_test.rs
@@ -33,7 +33,7 @@ fn test_simple_scenario() {
     scenario.run(false)
 }
 
-const PAGE_SIZE: usize = (size_of::<*const ()>() * 3).next_power_of_two();
+const PAGE_SIZE: usize = 64;
 const PAGE_SIZE_BITS: usize = PAGE_SIZE.trailing_zeros() as usize;
 const SIZE_CLASS_COUNT: usize = 4;
 
@@ -182,22 +182,17 @@ impl Scenario {
         let backing_mem = self
             .range_sizes
             .iter()
-            .map(|&size| unsafe {
-                mmap_anonymous(
-                    None,
-                    NonZero::new(size * PAGE_SIZE).unwrap(),
-                    ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
-                    MapFlags::MAP_PRIVATE,
-                )
-                .expect("failed to allocate memory")
+            .map(|&size| {
+                Mmap::new(NonZero::new(size * PAGE_SIZE).unwrap())
+                    .expect("failed to allocate memory")
             })
             .collect::<Vec<_>>();
 
         // Create the free list allocator and move the pages there.
         let mut free_list: FreeListAllocator<PAGE_SIZE> = FreeListAllocator::new();
-        for (&size, &ptr) in self.range_sizes.iter().zip(backing_mem.iter()) {
+        for (&size, mmap) in self.range_sizes.iter().zip(backing_mem.iter()) {
             unsafe {
-                free_list.add(ptr.cast(), size);
+                free_list.add(mmap.ptr(), size);
             }
         }
 
@@ -222,13 +217,6 @@ impl Scenario {
                 }
             }
         }
-
-        // Free each of the page ranges.
-        for (&size, ptr) in self.range_sizes.iter().zip(backing_mem) {
-            unsafe {
-                munmap(ptr, size * PAGE_SIZE).expect("failed to free memory");
-            }
-        }
     }
 }
 
@@ -256,3 +244,36 @@ impl<'alloc> fmt::Debug for Alloc<'alloc> {
             .finish()
     }
 }
+
+/// A mmap-allocated region of memory.
+struct Mmap {
+    ptr: NonNull<u8>,
+    len: NonZero<usize>,
+}
+
+impl Mmap {
+    pub fn new(len: NonZero<usize>) -> Result<Mmap, nix::Error> {
+        let ptr = unsafe {
+            mmap_anonymous(
+                None,
+                len,
+                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
+                MapFlags::MAP_PRIVATE,
+            )?
+            .cast()
+        };
+        Ok(Mmap { ptr, len })
+    }
+
+    pub fn ptr<T>(&self) -> NonNull<T> {
+        self.ptr.cast()
+    }
+}
+
+impl Drop for Mmap {
+    fn drop(&mut self) {
+        unsafe {
+            munmap(self.ptr(), self.len.get()).expect("failed to free memory");
+        }
+    }
+}