Add get_host_address_and_hugepages() to GuestMemory

Virtio-balloon can release the unused host memory to decrease the memory
usage of the VMM.
Release normal pages and hugepages require different operations.
(madvise MADV_DONTNEED and fallocate64 FALLOC_FL_PUNCH_HOLE)

This commit add get_host_address_and_hugepages() to GuestMemory to help
VMM decide if this is a hugepages address or not.

Signed-off-by: Hui Zhu <[email protected]>

Author: Hui Zhu

src/guest_memory.rs

@@ -288,6 +288,11 @@ pub trait GuestMemoryRegion: Bytes<MemoryRegionAddress, E = Error> {
     fn as_volatile_slice(&self) -> Result<volatile_memory::VolatileSlice> {
         self.get_slice(MemoryRegionAddress(0), self.len() as usize)
     }
+
+    /// Returns true if the region is hugepages
+    fn is_hugepages(&self) -> bool {
+        false
+    }
 }
 
 /// `GuestAddressSpace` provides a way to retrieve a `GuestMemory` object.
@@ -614,6 +619,13 @@ pub trait GuestMemory {
             .and_then(|(r, addr)| r.get_host_address(addr))
     }
 
+    /// Get the host virtual address corresponding to the guest address and the hugepages.
+    fn get_host_address_and_hugepages(&self, addr: GuestAddress) -> Result<(*mut u8, bool)> {
+        self.to_region_addr(addr)
+            .ok_or_else(|| Error::InvalidGuestAddress(addr))
+            .and_then(|(r, addr)| Ok((r.get_host_address(addr)?, r.is_hugepages())))
+    }
+
     /// Returns a [`VolatileSlice`](struct.VolatileSlice.html) of `count` bytes starting at
     /// `addr`.
     fn get_slice(

src/mmap.rs

@@ -403,6 +403,10 @@ impl GuestMemoryRegion for GuestRegionMmap {
         let slice = self.mapping.get_slice(offset.raw_value() as usize, count)?;
         Ok(slice)
     }
+
+    fn is_hugepages(&self) -> bool {
+        self.mapping.hugepages()
+    }
 }
 
 /// [`GuestMemory`](trait.GuestMemory.html) implementation that mmaps the guest's memory
@@ -434,6 +438,22 @@ impl GuestMemoryMmap {
     /// Valid memory regions are specified as a sequence of (Address, Size, Option<FileOffset>)
     /// tuples sorted by Address.
     pub fn from_ranges_with_files<A, T>(ranges: T) -> result::Result<Self, Error>
+    where
+        A: Borrow<(GuestAddress, usize, Option<FileOffset>)>,
+        T: IntoIterator<Item = A>,
+    {
+        Self::from_ranges_with_files_and_hugepages(ranges, false)
+    }
+
+    /// Creates a container and allocates anonymous memory for guest memory regions.
+    /// And set guest memory regions hugepages.
+    ///
+    /// Valid memory regions are specified as a sequence of (Address, Size, Option<FileOffset>)
+    /// tuples sorted by Address.
+    pub fn from_ranges_with_files_and_hugepages<A, T>(
+        ranges: T,
+        hugepages: bool,
+    ) -> result::Result<Self, Error>
     where
         A: Borrow<(GuestAddress, usize, Option<FileOffset>)>,
         T: IntoIterator<Item = A>,
@@ -445,13 +465,15 @@ impl GuestMemoryMmap {
                     let guest_base = x.borrow().0;
                     let size = x.borrow().1;
 
-                    if let Some(ref f_off) = x.borrow().2 {
+                    let mut r = if let Some(ref f_off) = x.borrow().2 {
                         MmapRegion::from_file(f_off.clone(), size)
                     } else {
                         MmapRegion::new(size)
                     }
-                    .map_err(Error::MmapRegion)
-                    .and_then(|r| GuestRegionMmap::new(r, guest_base))
+                    .map_err(Error::MmapRegion)?;
+                    r.set_hugepages(hugepages);
+
+                    GuestRegionMmap::new(r, guest_base)
                 })
                 .collect::<result::Result<Vec<_>, Error>>()?,
         )
@@ -989,6 +1011,64 @@ mod tests {
         }
     }
 
+    #[cfg(all(feature = "backend-mmap", unix))]
+    #[test]
+    fn test_get_host_address_and_hugepages() {
+        let f1 = TempFile::new().unwrap().into_file();
+        f1.set_len(0x400).unwrap();
+        let f2 = TempFile::new().unwrap().into_file();
+        f2.set_len(0x400).unwrap();
+
+        let start_addr1 = GuestAddress(0x0);
+        let start_addr2 = GuestAddress(0x800);
+        let guest_mem =
+            GuestMemoryMmap::from_ranges(&[(start_addr1, 0x400), (start_addr2, 0x400)]).unwrap();
+        let guest_mem_backed_by_file = GuestMemoryMmap::from_ranges_with_files_and_hugepages(
+            &[
+                (start_addr1, 0x400, Some(FileOffset::new(f1, 0))),
+                (start_addr2, 0x400, Some(FileOffset::new(f2, 0))),
+            ],
+            true,
+        )
+        .unwrap();
+
+        let guest_mem_list = vec![guest_mem];
+        for guest_mem in guest_mem_list.iter() {
+            assert!(guest_mem.get_host_address(GuestAddress(0x600)).is_err());
+            let (ptr0, ptr0_huge) = guest_mem
+                .get_host_address_and_hugepages(GuestAddress(0x800))
+                .unwrap();
+            let (ptr1, ptr1_huge) = guest_mem
+                .get_host_address_and_hugepages(GuestAddress(0xa00))
+                .unwrap();
+            assert_eq!(
+                ptr0,
+                guest_mem.find_region(GuestAddress(0x800)).unwrap().as_ptr()
+            );
+            assert_eq!(ptr0_huge, false);
+            assert_eq!(ptr1_huge, false);
+            assert_eq!(unsafe { ptr0.offset(0x200) }, ptr1);
+        }
+
+        let guest_mem_list = vec![guest_mem_backed_by_file];
+        for guest_mem in guest_mem_list.iter() {
+            assert!(guest_mem.get_host_address(GuestAddress(0x600)).is_err());
+            let (ptr0, ptr0_huge) = guest_mem
+                .get_host_address_and_hugepages(GuestAddress(0x800))
+                .unwrap();
+            let (ptr1, ptr1_huge) = guest_mem
+                .get_host_address_and_hugepages(GuestAddress(0xa00))
+                .unwrap();
+            assert_eq!(
+                ptr0,
+                guest_mem.find_region(GuestAddress(0x800)).unwrap().as_ptr()
+            );
+            assert_eq!(ptr0_huge, true);
+            assert_eq!(ptr1_huge, true);
+            assert_eq!(unsafe { ptr0.offset(0x200) }, ptr1);
+        }
+    }
+
     #[test]
     fn test_deref() {
         let f = TempFile::new().unwrap().into_file();

src/mmap_unix.rs

@@ -90,6 +90,7 @@ pub struct MmapRegion {
     prot: i32,
     flags: i32,
     owned: bool,
+    hugepages: bool,
 }
 
 // Send and Sync aren't automatically inherited for the raw address pointer.
@@ -172,6 +173,7 @@ impl MmapRegion {
             prot,
             flags,
             owned: true,
+            hugepages: false,
         })
     }
 
@@ -213,6 +215,7 @@ impl MmapRegion {
             prot,
             flags,
             owned: false,
+            hugepages: false,
         })
     }
 
@@ -248,6 +251,16 @@ impl MmapRegion {
         self.owned
     }
 
+    /// Set the hugepages of the region
+    pub fn set_hugepages(&mut self, hugepages: bool) {
+        self.hugepages = hugepages
+    }
+
+    /// Returns `true` if the region is hugepages
+    pub fn hugepages(&self) -> bool {
+        self.hugepages
+    }
+
     /// Checks whether this region and `other` are backed by overlapping
     /// [`FileOffset`](struct.FileOffset.html) objects.
     ///

src/mmap_windows.rs

@@ -175,6 +175,12 @@ impl MmapRegion {
     pub fn file_offset(&self) -> Option<&FileOffset> {
         self.file_offset.as_ref()
     }
+
+    pub fn set_hugepages(&mut self, _hugepages: bool) {}
+
+    pub fn hugepages(&self) -> bool {
+        false
+    }
 }
 
 impl AsSlice for MmapRegion {