fix(memory): punch holes for shared discard ranges

Use fallocate(PUNCH_HOLE|KEEP_SIZE) for MAP_SHARED file-backed guest
memory so memfd-backed balloon hinting/reporting clears the shared
backing instead of only dropping PTEs.

Signed-off-by: Babis Chalios <babis.chalios@e2b.dev>

Author: ValentaTomas

src/vmm/src/test_utils/mod.rs

@@ -3,6 +3,7 @@
 
 #![allow(missing_docs)]
 
+use std::fs::File;
 use std::sync::{Arc, Mutex};
 
 use vm_memory::{GuestAddress, GuestRegionCollection};
@@ -53,6 +54,22 @@ pub fn multi_region_mem(regions: &[(GuestAddress, usize)]) -> GuestMemoryMmap {
     .unwrap()
 }
 
+/// Creates a [`GuestMemoryMmap`] with multiple regions and without dirty page tracking.
+pub fn multi_region_mem_memfd(
+    regions: &[(GuestAddress, usize)],
+    huge_page_cfg: HugePageConfig,
+) -> (GuestMemoryMmap, Arc<File>) {
+    let (reg, file) = memory::memfd_backed(regions, false, huge_page_cfg).unwrap();
+    let mem = GuestRegionCollection::from_regions(
+        reg.into_iter()
+            .map(|region| GuestRegionMmapExt::dram_from_mmap_region(region, 0))
+            .collect(),
+    )
+    .unwrap();
+
+    (mem, file)
+}
+
 pub fn multi_region_mem_raw(regions: &[(GuestAddress, usize)]) -> Vec<GuestRegionMmap> {
     memory::anonymous(regions.iter().copied(), false, HugePageConfig::None)
         .expect("Cannot initialize memory")

src/vmm/src/vstate/memory.rs

@@ -22,9 +22,10 @@ pub use vm_memory::{
     GuestUsize, MemoryRegionAddress, MmapRegion, address,
 };
 use vm_memory::{GuestMemoryError, GuestMemoryRegionBytes, VolatileSlice, WriteVolatile};
-use vmm_sys_util::errno;
+use vmm_sys_util::fallocate::FallocateMode;
+use vmm_sys_util::{errno, fallocate};
 
-use crate::utils::{get_page_size, u64_to_usize};
+use crate::utils::{get_page_size, u64_to_usize, usize_to_u64};
 use crate::vmm_config::machine_config::HugePageConfig;
 use crate::vstate::vm::VmError;
 use crate::{DirtyBitmap, Vm};
@@ -391,8 +392,7 @@ impl GuestRegionMmapExt {
         let phys_address = self.get_host_address(caddr)?;
 
         match (self.inner.file_offset(), self.inner.flags()) {
-            // If and only if we are resuming from a snapshot file, we have a file and it's mapped
-            // private
+            // If we are resuming from a snapshot file, we have a file and it's mapped private
             (Some(_), flags) if flags & libc::MAP_PRIVATE != 0 => {
                 // Mmap a new anonymous region over the present one in order to create a hole
                 // with zero pages.
@@ -420,12 +420,48 @@ impl GuestRegionMmapExt {
                     Ok(())
                 }
             }
-            // Match either the case of an anonymous mapping, or the case
-            // of a shared file mapping.
-            // TODO: madvise(MADV_DONTNEED) doesn't actually work with memfd
-            // (or in general MAP_SHARED of a fd). In those cases we should use
-            // fallocate64(FALLOC_FL_PUNCH_HOLE|FALLOC_FL_KEEP_SIZE).
-            // We keep falling to the madvise branch to keep the previous behaviour.
+            // If we back memory over memfd we have a file mapped shared.
+            (Some(file_offset), flags) if flags & libc::MAP_SHARED != 0 => {
+                let Some(offset) = file_offset.start().checked_add(caddr.raw_value()) else {
+                    return Err(GuestMemoryError::InvalidGuestAddress(GuestAddress(
+                        caddr.raw_value(),
+                    )));
+                };
+
+                // fallocate(PUNCH_HOLE) silently no-ops on memfd when [offset, offset+len)
+                // doesn't cover any whole filesystem block, and returns EINVAL on hugetlbfs
+                // for ranges not aligned to the huge-page size. Expand [offset, offset+len)
+                // to the enclosing page boundaries so the punch actually covers every page
+                // the caller asked to discard.
+                let page_size = usize_to_u64(if flags & libc::MAP_HUGETLB != 0 {
+                    HugePageConfig::Hugetlbfs2M.page_size()
+                } else {
+                    get_page_size().map_err(|err| GuestMemoryError::IOError(err.into()))?
+                });
+                let end = offset
+                    .checked_add(usize_to_u64(len))
+                    .ok_or(GuestMemoryError::GuestAddressOverflow)?;
+                let aligned_offset = offset & !(page_size - 1);
+                let aligned_end = end
+                    .checked_add(page_size - 1)
+                    .ok_or(GuestMemoryError::GuestAddressOverflow)?
+                    & !(page_size - 1);
+
+                fallocate::fallocate(
+                    file_offset.file(),
+                    FallocateMode::PunchHole,
+                    true,
+                    aligned_offset,
+                    aligned_end - aligned_offset,
+                )
+                .map_err(|err| {
+                    error!("discard_range: punching hole failed: {err:?}");
+                    GuestMemoryError::IOError(err.into())
+                })?;
+
+                Ok(())
+            }
+            // Anonymous mapping.
             _ => {
                 // Madvise the region in order to mark it as not used.
                 // SAFETY: The address and length are known to be valid.
@@ -873,7 +909,7 @@ mod tests {
 
     use super::*;
     use crate::snapshot::Snapshot;
-    use crate::test_utils::single_region_mem;
+    use crate::test_utils::{multi_region_mem_memfd, single_region_mem};
     use crate::utils::{get_page_size, mib_to_bytes};
     use crate::vstate::memory::test_utils::into_region_ext;
 
@@ -1462,6 +1498,92 @@ mod tests {
         );
     }
 
+    fn check_pages(
+        mem: &GuestMemoryMmap,
+        page_size: usize,
+        first_page_expected: &[u8],
+        second_page_expected: &[u8],
+        third_page_expected: &[u8],
+    ) {
+        let first_page = 0u64;
+        let second_page = usize_to_u64(page_size);
+        let third_page = usize_to_u64(2 * page_size);
+
+        let mut actual_page = vec![0u8; page_size];
+        mem.read(actual_page.as_mut_slice(), GuestAddress(first_page))
+            .unwrap();
+        assert_eq!(first_page_expected, actual_page);
+        mem.read(actual_page.as_mut_slice(), GuestAddress(second_page))
+            .unwrap();
+        assert_eq!(second_page_expected, actual_page);
+        mem.read(actual_page.as_mut_slice(), GuestAddress(third_page))
+            .unwrap();
+        assert_eq!(third_page_expected, actual_page);
+    }
+
+    // A page size-agnostic tests for discard_range() of memory that is
+    // mapped on top of memfd
+    fn test_discard_range_on_shared_memfd(huge_pages: HugePageConfig) {
+        let page_size = huge_pages.page_size();
+        let first_page = 0u64;
+        let second_page = usize_to_u64(page_size);
+        let third_page = usize_to_u64(2 * page_size);
+
+        let (mem, _file) = multi_region_mem_memfd(
+            &[
+                (GuestAddress(first_page), page_size),
+                (GuestAddress(second_page), page_size),
+                (GuestAddress(third_page), page_size),
+            ],
+            huge_pages,
+        );
+
+        // Fill up all three pages with 1s
+        let ones = vec![1u8; 3 * page_size];
+        mem.write(&ones, GuestAddress(0)).unwrap();
+
+        // Cleanup the last page
+        mem.discard_range(GuestAddress(third_page), page_size)
+            .unwrap();
+        check_pages(
+            &mem,
+            page_size,
+            &vec![1u8; page_size],
+            &vec![1u8; page_size],
+            &vec![0u8; page_size],
+        );
+
+        // Cleanup the first page
+        mem.discard_range(GuestAddress(0), page_size).unwrap();
+        check_pages(
+            &mem,
+            page_size,
+            &vec![0u8; page_size],
+            &vec![1u8; page_size],
+            &vec![0u8; page_size],
+        );
+
+        // Cleanup starting from first page and landing in second
+        mem.discard_range(GuestAddress(1), page_size).unwrap();
+        check_pages(
+            &mem,
+            page_size,
+            &vec![0u8; page_size],
+            &vec![0u8; page_size],
+            &vec![0u8; page_size],
+        );
+    }
+
+    #[test]
+    fn test_discard_range_on_shared_memfd_4k() {
+        test_discard_range_on_shared_memfd(HugePageConfig::None)
+    }
+
+    #[test]
+    fn test_discard_range_on_shared_memfd_2m() {
+        test_discard_range_on_shared_memfd(HugePageConfig::Hugetlbfs2M)
+    }
+
     /// Verifies that `slots_intersecting_range` returns the correct slots for
     /// ranges at slot boundaries, interior to a slot, and spanning two slots.
     #[test]