feat(balloon): add support for discarding memfd-backed memory

In case the backing is a fd that's mapped shared (eg memfd), we need to
fallocate in order to drop the memory page.

Signed-off-by: Riccardo Mancini <[email protected]>

Author: Manciukic

src/vmm/src/vstate/memory.rs

@@ -7,6 +7,7 @@
 
 use std::fs::File;
 use std::io::SeekFrom;
+use std::os::fd::AsRawFd;
 use std::sync::Arc;
 
 use kvm_bindings::{KVM_MEM_LOG_DIRTY_PAGES, kvm_userspace_memory_region};
@@ -119,6 +120,7 @@ impl GuestMemoryRegionExt {
                 // file only drops any anonymous pages in range, but subsequent accesses would read
                 // whatever page is stored on the backing file. Mmapping anonymous pages ensures
                 // it's zeroed.
+                // Punching a hole doesn't work in this case as the file is read-only.
                 // SAFETY: The address and length are known to be valid.
                 let ret = unsafe {
                     libc::mmap(
@@ -138,9 +140,33 @@ impl GuestMemoryRegionExt {
                     Ok(len)
                 }
             }
-            // TODO: memfd doesn't actually work with madvise, but this is to keep the
-            // previous behaviour
-            (None, _) | (Some(_), _) => {
+            // If and only if we are using fd-backed memory (eg memfd for vhost-user), we have a
+            // file descriptor and it's mapped shared.
+            (Some(file_offset), flags) if flags & libc::MAP_SHARED != 0 => {
+                // Fallocate to punch a hole in the region to free memory pages.
+                // In case of a FD mmapped as shared (eg memfd), MADV_DONTNEED doesn't work
+                // and we can't just mmap a zero page on top of it as that doesn't actually
+                // change the underlying shared FD, but we can punch a hole into it.
+                // SAFETY: The address and length are known to be valid.
+                let ret = unsafe {
+                    libc::fallocate64(
+                        file_offset.file().as_raw_fd(),
+                        libc::FALLOC_FL_PUNCH_HOLE | libc::FALLOC_FL_KEEP_SIZE,
+                        i64::try_from(file_offset.start() + caddr.0).unwrap(),
+                        len.try_into().unwrap(),
+                    )
+                };
+                if ret < 0 {
+                    let os_error = std::io::Error::last_os_error();
+                    error!("discard_range: fallocate64 failed: {:?}", os_error);
+                    Err(GuestMemoryError::IOError(os_error))
+                } else {
+                    Ok(len)
+                }
+            }
+            // Either it's mapped PRIVATE or SHARED
+            (Some(_), _) => unreachable!(),
+            (None, _) => {
                 // Madvise the region in order to mark it as not used.
                 // SAFETY: The address and length are known to be valid.
                 let ret = unsafe { libc::madvise(phys_address.cast(), len, libc::MADV_DONTNEED) };
@@ -1044,4 +1070,49 @@ mod tests {
             GuestMemoryError::IOError(_)
         );
     }
+
+    #[test]
+    fn test_discard_range_on_memfd() {
+        let page_size: usize = 0x1000;
+        let mem = into_region_ext(
+            memfd_backed(
+                &[(GuestAddress(0), 2 * page_size)],
+                false,
+                HugePageConfig::None,
+            )
+            .unwrap(),
+        );
+
+        // Fill the memory with ones.
+        let ones = vec![1u8; 2 * page_size];
+        mem.write(&ones[..], GuestAddress(0)).unwrap();
+
+        // Remove the first page.
+        mem.discard_range(GuestAddress(0), page_size).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!(
+            mem.discard_range(GuestAddress(0), 0x10000).unwrap_err(),
+            GuestMemoryError::PartialBuffer {
+                expected: _,
+                completed: _
+            }
+        );
+
+        // Region not mapped.
+        assert_match!(
+            mem.discard_range(GuestAddress(0x10000), 0x10).unwrap_err(),
+            GuestMemoryError::InvalidGuestAddress(_)
+        );
+    }
 }