[RFC] Use guest memory collection and other refactors

src/vmm/src/devices/virtio/balloon/device.rs

@@ -14,7 +14,7 @@ use super::super::ActivateError;
 use super::super::device::{DeviceState, VirtioDevice};
 use super::super::queue::Queue;
 use super::metrics::METRICS;
-use super::util::{compact_page_frame_numbers, remove_range};
+use super::util::compact_page_frame_numbers;
 use super::{
     BALLOON_DEV_ID, BALLOON_NUM_QUEUES, BALLOON_QUEUE_SIZES, DEFLATE_INDEX, INFLATE_INDEX,
     MAX_PAGE_COMPACT_BUFFER, MAX_PAGES_IN_DESC, MIB_TO_4K_PAGES, STATS_INDEX,
@@ -32,7 +32,9 @@ use crate::devices::virtio::queue::InvalidAvailIdx;
 use crate::devices::virtio::transport::{VirtioInterrupt, VirtioInterruptType};
 use crate::logger::IncMetric;
 use crate::utils::u64_to_usize;
-use crate::vstate::memory::{Address, ByteValued, Bytes, GuestAddress, GuestMemoryMmap};
+use crate::vstate::memory::{
+    Address, ByteValued, Bytes, GuestAddress, GuestMemoryExtension, GuestMemoryMmap,
+};
 use crate::{impl_device_type, mem_size_mib};
 
 const SIZE_OF_U32: usize = std::mem::size_of::<u32>();
@@ -335,10 +337,9 @@ impl Balloon {
                 let guest_addr =
                     GuestAddress(u64::from(page_frame_number) << VIRTIO_BALLOON_PFN_SHIFT);
 
-                if let Err(err) = remove_range(
-                    mem,
-                    (guest_addr, u64::from(range_len) << VIRTIO_BALLOON_PFN_SHIFT),
-                    self.restored_from_file,
+                if let Err(err) = mem.discard_range(
+                    guest_addr,
+                    usize::try_from(range_len).unwrap() << VIRTIO_BALLOON_PFN_SHIFT,
                 ) {
                     error!("Error removing memory range: {:?}", err);
                 }

src/vmm/src/devices/virtio/balloon/util.rs

@@ -65,57 +65,6 @@ pub(crate) fn compact_page_frame_numbers(v: &mut [u32]) -> Vec<(u32, u32)> {
     result
 }
 
-pub(crate) fn remove_range(
-    guest_memory: &GuestMemoryMmap,
-    range: (GuestAddress, u64),
-    restored_from_file: bool,
-) -> Result<(), RemoveRegionError> {
-    let (guest_address, range_len) = range;
-
-    if let Some(region) = guest_memory.find_region(guest_address) {
-        if guest_address.0 + range_len > region.start_addr().0 + region.len() {
-            return Err(RemoveRegionError::MalformedRange);
-        }
-        let phys_address = guest_memory
-            .get_host_address(guest_address)
-            .map_err(|_| RemoveRegionError::AddressTranslation)?;
-
-        // Mmap a new anonymous region over the present one in order to create a hole.
-        // This workaround is (only) needed after resuming from a snapshot because the guest memory
-        // is mmaped from file as private and there is no `madvise` flag that works for this case.
-        if restored_from_file {
-            // SAFETY: The address and length are known to be valid.
-            let ret = unsafe {
-                libc::mmap(
-                    phys_address.cast(),
-                    u64_to_usize(range_len),
-                    libc::PROT_READ | libc::PROT_WRITE,
-                    libc::MAP_FIXED | libc::MAP_ANONYMOUS | libc::MAP_PRIVATE,
-                    -1,
-                    0,
-                )
-            };
-            if ret == libc::MAP_FAILED {
-                return Err(RemoveRegionError::MmapFail(io::Error::last_os_error()));
-            }
-        };
-
-        // Madvise the region in order to mark it as not used.
-        // SAFETY: The address and length are known to be valid.
-        let ret = unsafe {
-            let range_len = u64_to_usize(range_len);
-            libc::madvise(phys_address.cast(), range_len, libc::MADV_DONTNEED)
-        };
-        if ret < 0 {
-            return Err(RemoveRegionError::MadviseFail(io::Error::last_os_error()));
-        }
-
-        Ok(())
-    } else {
-        Err(RemoveRegionError::RegionNotFound)
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use std::fmt::Debug;
@@ -174,88 +123,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn test_remove_range() {
-        let page_size: usize = 0x1000;
-        let mem = single_region_mem(2 * page_size);
-
-        // Fill the memory with ones.
-        let ones = vec![1u8; 2 * page_size];
-        mem.write(&ones[..], GuestAddress(0)).unwrap();
-
-        // Remove the first page.
-        remove_range(&mem, (GuestAddress(0), page_size as u64), false).unwrap();
-
-        // Check that the first page is zeroed.
-        let mut actual_page = vec![0u8; page_size];
-        mem.read(actual_page.as_mut_slice(), GuestAddress(0))
-            .unwrap();
-        assert_eq!(vec![0u8; page_size], actual_page);
-        // Check that the second page still contains ones.
-        mem.read(actual_page.as_mut_slice(), GuestAddress(page_size as u64))
-            .unwrap();
-        assert_eq!(vec![1u8; page_size], actual_page);
-
-        // Malformed range: the len is too big.
-        assert_match!(
-            remove_range(&mem, (GuestAddress(0), 0x10000), false).unwrap_err(),
-            RemoveRegionError::MalformedRange
-        );
-
-        // Region not mapped.
-        assert_match!(
-            remove_range(&mem, (GuestAddress(0x10000), 0x10), false).unwrap_err(),
-            RemoveRegionError::RegionNotFound
-        );
-
-        // Madvise fail: the guest address is not aligned to the page size.
-        assert_match!(
-            remove_range(&mem, (GuestAddress(0x20), page_size as u64), false).unwrap_err(),
-            RemoveRegionError::MadviseFail(_)
-        );
-    }
-
-    #[test]
-    fn test_remove_range_on_restored() {
-        let page_size: usize = 0x1000;
-        let mem = single_region_mem(2 * page_size);
-
-        // Fill the memory with ones.
-        let ones = vec![1u8; 2 * page_size];
-        mem.write(&ones[..], GuestAddress(0)).unwrap();
-
-        // Remove the first page.
-        remove_range(&mem, (GuestAddress(0), page_size as u64), true).unwrap();
-
-        // Check that the first page is zeroed.
-        let mut actual_page = vec![0u8; page_size];
-        mem.read(actual_page.as_mut_slice(), GuestAddress(0))
-            .unwrap();
-        assert_eq!(vec![0u8; page_size], actual_page);
-        // Check that the second page still contains ones.
-        mem.read(actual_page.as_mut_slice(), GuestAddress(page_size as u64))
-            .unwrap();
-        assert_eq!(vec![1u8; page_size], actual_page);
-
-        // Malformed range: the len is too big.
-        assert_match!(
-            remove_range(&mem, (GuestAddress(0), 0x10000), true).unwrap_err(),
-            RemoveRegionError::MalformedRange
-        );
-
-        // Region not mapped.
-        assert_match!(
-            remove_range(&mem, (GuestAddress(0x10000), 0x10), true).unwrap_err(),
-            RemoveRegionError::RegionNotFound
-        );
-
-        // Mmap fail: the guest address is not aligned to the page size.
-        assert_match!(
-            remove_range(&mem, (GuestAddress(0x20), page_size as u64), true).unwrap_err(),
-            RemoveRegionError::MmapFail(_)
-        );
-    }
-
     /// -------------------------------------
     /// BEGIN PROPERTY BASED TESTING
     use proptest::prelude::*;

src/vmm/src/test_utils/mod.rs

@@ -16,8 +16,7 @@ use crate::vm_memory_vendored::GuestRegionCollection;
 use crate::vmm_config::boot_source::BootSourceConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::HugePageConfig;
-use crate::vstate::memory;
-use crate::vstate::memory::{GuestMemoryMmap, GuestRegionMmap, GuestRegionMmapExt};
+use crate::vstate::memory::{self, GuestMemoryMmap, GuestRegionMmap, GuestRegionMmapExt};
 use crate::{EventManager, Vmm};
 
 pub mod mock_resources;