refactor(balloon): move remove_range to GuestMemoryExtension

Move the function to discard guest memory under the GuestMemoryExtension
trait.
This is done so that it could be reused in other parts of the code.

Furthermore, this also changes it to use try_access, which is the
correct way to access a range of guest memory as it deals with the
possibility of it spanning multiple regions.

Signed-off-by: Riccardo Mancini <mancio@amazon.com>

Author: Manciukic

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

@@ -15,8 +15,7 @@ use crate::test_utils::mock_resources::{MockBootSourceConfig, MockVmConfig, Mock
 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;